Netpulse_SasS/agent/internal/ncmx/session_test.go

package ncmx

import (
	"bytes"
	"context"
	"io"
	"regexp"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakeDevice — пристрій, що відповідає за сценарієм.
//
// Саме пристрій, а не мок з'єднання: перевіряти треба поведінку розбору
// на справжньому потоці байтів, включно з тим, що промпт приходить
// розірваним між читаннями.
type fakeDevice struct {
	mu       sync.Mutex
	out      bytes.Buffer
	replies  map[string]string
	banner   string
	prompt   string
	written  []string
	chunkLen int
	closed   bool
}

func newDevice(prompt, banner string, replies map[string]string) *fakeDevice {
	d := &fakeDevice{replies: replies, banner: banner, prompt: prompt}
	d.out.WriteString(banner + prompt)
	return d
}

func (d *fakeDevice) Read(p []byte) (int, error) {
	deadline := time.Now().Add(2 * time.Second)
	for {
		d.mu.Lock()
		if d.out.Len() > 0 {
			n := d.out.Len()
			// Віддаємо дрібними шматками: так промпт гарантовано
			// розривається між викликами Read, і код мусить це пережити.
			if d.chunkLen > 0 && n > d.chunkLen {
				n = d.chunkLen
			}
			if n > len(p) {
				n = len(p)
			}
			nn, _ := d.out.Read(p[:n])
			d.mu.Unlock()
			return nn, nil
		}
		closed := d.closed
		d.mu.Unlock()
		if closed {
			return 0, io.EOF
		}
		if time.Now().After(deadline) {
			return 0, io.EOF
		}
		time.Sleep(time.Millisecond)
	}
}

func (d *fakeDevice) Write(p []byte) (int, error) {
	cmd := strings.TrimRight(string(p), "\r\n")
	d.mu.Lock()
	defer d.mu.Unlock()
	d.written = append(d.written, cmd)

	// Відлуння команди — так поводиться справжня консоль.
	d.out.WriteString(cmd + "\r\n")
	if reply, ok := d.replies[cmd]; ok {
		d.out.WriteString(reply)
	}
	d.out.WriteString(d.prompt)
	return len(p), nil
}

func (d *fakeDevice) Close() error {
	d.mu.Lock()
	d.closed = true
	d.mu.Unlock()
	return nil
}

func cliFor(d *fakeDevice, prompt string) *CLI {
	return NewCLI(d, Options{
		PromptRe:    regexp.MustCompile(`(?m)` + prompt),
		IdleTimeout: time.Second,
	})
}

func TestRunStripsEchoAndPrompt(t *testing.T) {
	d := newDevice("\nsw1#", "Welcome\r\n", map[string]string{
		"show running-config": "hostname sw1\r\ninterface Gi0/1\r\n",
	})
	c := cliFor(d, `[>#]\s*$`)
	ctx := context.Background()

	if err := c.WaitPrompt(ctx); err != nil {
		t.Fatalf("привітання: %v", err)
	}

	out, err := c.Run(ctx, "show running-config")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if out != "hostname sw1\ninterface Gi0/1" {
		t.Fatalf("отримали %q", out)
	}
}

// Промпт, розірваний між читаннями, — норма для повільного каналу.
func TestRunHandlesSplitPrompt(t *testing.T) {
	d := newDevice("\nsw1#", "", map[string]string{
		"show run": "line one\r\nline two\r\n",
	})
	d.chunkLen = 3 // рвемо потік по три байти

	c := cliFor(d, `[>#]\s*$`)
	ctx := context.Background()
	_ = c.WaitPrompt(ctx)

	out, err := c.Run(ctx, "show run")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if out != "line one\nline two" {
		t.Fatalf("отримали %q", out)
	}
}

// Рядок конфігу, схожий на запрошення, не має обривати збір.
//
// Це не теоретичний випадок: banner motd із «#» усередині трапляється
// в кожній другій мережі.
func TestPromptInsideBodyDoesNotTruncate(t *testing.T) {
	body := "banner motd #\r\nDANGER\r\n#\r\nhostname sw1\r\n"
	d := newDevice("\nsw1#", "", map[string]string{"show run": body})

	c := cliFor(d, `sw1[>#]\s*$`)
	ctx := context.Background()
	_ = c.WaitPrompt(ctx)

	out, err := c.Run(ctx, "show run")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if !strings.Contains(out, "hostname sw1") {
		t.Fatalf("тіло обрізане: %q", out)
	}
}

// Посторінковий вивід: агент має відповісти й прибрати «--More--» із тексту.
//
// Макет тут керується вручну, а не через таблицю відповідей: справжній
// пристрій, зупинившись на «--More--», НЕ шле запрошення — саме тому
// агент і мусить його розпізнати. Автоматичне дописування промпту
// зробило б перевірку беззмістовною.
func TestMorePromptIsAnsweredAndStripped(t *testing.T) {
	d := &scriptedDevice{prompt: "\nsw1#"}
	d.push("\nsw1#") // початкове запрошення

	c := NewCLI(d, Options{
		PromptRe:    regexp.MustCompile(`(?m)[>#]\s*$`),
		IdleTimeout: 2 * time.Second,
	})
	ctx := context.Background()
	if err := c.WaitPrompt(ctx); err != nil {
		t.Fatalf("привітання: %v", err)
	}

	// Пристрій віддає першу сторінку й зупиняється на пейджері.
	d.onWrite("show run", "show run\r\npart one\r\n --More-- ")
	// На пробіл — решта й нарешті запрошення.
	d.onWrite(" ", "\rpart two\r\nsw1#")

	out, err := c.Run(ctx, "show run")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if strings.Contains(out, "More") {
		t.Fatalf("запрошення пейджера потрапило в конфіг: %q", out)
	}
	if !strings.Contains(out, "part one") || !strings.Contains(out, "part two") {
		t.Fatalf("вивід зібрано не повністю: %q", out)
	}
	if got := d.sent(); len(got) < 2 || got[1] != " " {
		t.Fatalf("агент не відповів пробілом на пейджер: %v", got)
	}
}

// scriptedDevice відповідає рівно тим, що йому прописали на конкретний
// запис. Нічого не додає від себе — на відміну від fakeDevice, який
// імітує звичайну консоль.
type scriptedDevice struct {
	mu      sync.Mutex
	out     bytes.Buffer
	script  map[string]string
	written []string
	prompt  string
	closed  bool
}

func (d *scriptedDevice) push(s string) {
	d.mu.Lock()
	d.out.WriteString(s)
	d.mu.Unlock()
}

func (d *scriptedDevice) onWrite(in, reply string) {
	d.mu.Lock()
	if d.script == nil {
		d.script = map[string]string{}
	}
	d.script[in] = reply
	d.mu.Unlock()
}

func (d *scriptedDevice) sent() []string {
	d.mu.Lock()
	defer d.mu.Unlock()
	return append([]string(nil), d.written...)
}

func (d *scriptedDevice) Read(p []byte) (int, error) {
	deadline := time.Now().Add(2 * time.Second)
	for {
		d.mu.Lock()
		if d.out.Len() > 0 {
			n, _ := d.out.Read(p)
			d.mu.Unlock()
			return n, nil
		}
		closed := d.closed
		d.mu.Unlock()
		if closed || time.Now().After(deadline) {
			return 0, io.EOF
		}
		time.Sleep(time.Millisecond)
	}
}

func (d *scriptedDevice) Write(p []byte) (int, error) {
	// Пейджеру відповідають голим пробілом без переводу рядка, тому
	// ключем є те, що надіслали, як є.
	key := strings.TrimRight(string(p), "\r\n")
	if key == "" {
		key = string(p)
	}

	d.mu.Lock()
	d.written = append(d.written, key)
	if reply, ok := d.script[key]; ok {
		d.out.WriteString(reply)
	}
	d.mu.Unlock()
	return len(p), nil
}

func (d *scriptedDevice) Close() error {
	d.mu.Lock()
	d.closed = true
	d.mu.Unlock()
	return nil
}

// Тиша в каналі має давати саме ErrPromptTimeout: це підказує, що
// проблема в prompt_regex, а не в мережі.
func TestSilenceGivesPromptTimeout(t *testing.T) {
	d := newDevice("\nsw1#", "", nil)
	c := NewCLI(d, Options{
		PromptRe:    regexp.MustCompile(`(?m)ЦЬОГО-НЕ-БУДЕ$`),
		IdleTimeout: 150 * time.Millisecond,
	})

	_, err := c.readUntilPrompt(context.Background())
	if err != ErrPromptTimeout {
		t.Fatalf("очікували ErrPromptTimeout, отримали %v", err)
	}
}

func TestMaxBytesGuard(t *testing.T) {
	huge := strings.Repeat("x", 40000) + "\r\n"
	d := newDevice("\nsw1#", "", map[string]string{"show run": huge})

	c := NewCLI(d, Options{
		PromptRe:    regexp.MustCompile(`(?m)[>#]\s*$`),
		IdleTimeout: time.Second,
		MaxBytes:    1024,
	})
	ctx := context.Background()
	_ = c.WaitPrompt(ctx)

	if _, err := c.Run(ctx, "show run"); err != ErrTooLarge {
		t.Fatalf("очікували ErrTooLarge, отримали %v", err)
	}
}

// Пристрій без відлуння не має втрачати перший рядок конфігу.
func TestNoEchoKeepsFirstLine(t *testing.T) {
	raw := "hostname sw1\ninterface Gi0/1"
	if got := cleanOutput(raw, "show run"); got != raw {
		t.Fatalf("перший рядок з'їдено: %q", got)
	}
}

func TestEchoIsRemovedOnce(t *testing.T) {
	raw := "show run\nhostname sw1\nshow run\n"
	got := cleanOutput(raw, "show run")
	if got != "hostname sw1\nshow run" {
		t.Fatalf("отримали %q", got)
	}
}

func TestStripFirstLines(t *testing.T) {
	s := "Building configuration...\n\nCurrent configuration : 100 bytes\nhostname sw1"
	if got := StripFirstLines(s, 3); got != "hostname sw1" {
		t.Fatalf("отримали %q", got)
	}
	if got := StripFirstLines(s, 0); got != s {
		t.Fatal("нуль рядків має лишати текст як є")
	}
	if got := StripFirstLines("одне", 5); got != "" {
		t.Fatalf("забагато рядків: %q", got)
	}
}

// Порожній prompt_regex не має валити збір: підставляємо типовий.
func TestEmptyPromptFallsBack(t *testing.T) {
	re, err := compilePrompt("")
	if err != nil {
		t.Fatalf("compilePrompt: %v", err)
	}
	if !re.MatchString("sw1#") {
		t.Fatal("типовий шаблон не ловить «sw1#»")
	}
}

func TestBadPromptIsReported(t *testing.T) {
	if _, err := compilePrompt("[unclosed"); err == nil {
		t.Fatal("некоректний регекс мав дати помилку")
	}
}

// Багаторядковий режим обов'язковий: без нього «$» шукає кінець усього
// тексту, а запрошення стоїть у кінці останнього рядка.
func TestPromptIsMultiline(t *testing.T) {
	re, _ := compilePrompt(`[>#]\s*$`)
	if !re.MatchString("рядок\nще рядок\nsw1#") {
		t.Fatal("шаблон не працює в багаторядковому тексті")
	}
}

// Керуючі коди терміналу не мають потрапляти в конфіг: перший же diff
// показав би зміну там, де змінився лише режим терміналу.
func TestAnsiSequencesAreStripped(t *testing.T) {
	raw := "\x1b[?2004lhostname sw1\x1b[0m\r\ninterface Gi0/1\r\n"
	got := cleanOutput(raw, "show run")
	if strings.ContainsRune(got, 0x1b) {
		t.Fatalf("escape-послідовність лишилась: %q", got)
	}
	if got != "hostname sw1\ninterface Gi0/1" {
		t.Fatalf("отримали %q", got)
	}
}

// Псевдотермінал видає «\r\r\n»: сам пристрій шле «\r\n», PTY додає свій
// «\r». Без згортання кожен рядок подвоювався б.
func TestDoubleCarriageReturnDoesNotDoubleLines(t *testing.T) {
	raw := "line one\r\r\nline two\r\r\n"
	got := cleanOutput(raw, "show run")
	if got != "line one\nline two" {
		t.Fatalf("отримали %q (рядків %d)", got, strings.Count(got, "\n")+1)
	}
}

// Відлуння команди трапляється не першим рядком, а після залишку
// запрошення — саме так поводиться bash у PTY.
func TestEchoAfterPromptRemnantIsRemoved(t *testing.T) {
	raw := "claude@host:~$ cat /etc/os-release\r\nID=debian\r\n"
	got := cleanOutput(raw, "cat /etc/os-release")
	if strings.Contains(got, "os-release") {
		t.Fatalf("відлуння лишилось: %q", got)
	}
	if got != "ID=debian" {
		t.Fatalf("отримали %q", got)
	}
}