fetch_ml/tests/fixtures/tui_driver.go

// Package tests provides TUI test driver for SSH testing
package tests

import (
	"bytes"
	"fmt"
	"io"
	"time"

	"golang.org/x/crypto/ssh"
)

// TUIDriver provides a driver for interacting with the TUI over SSH
type TUIDriver struct {
	session   *ssh.Session
	stdin     io.WriteCloser
	stdout    io.Reader
	stderr    io.Reader
	outputBuf *bytes.Buffer
	width     int
	height    int
}

// NewTUIDriver creates a new TUI driver connected via SSH
func NewTUIDriver(session *ssh.Session, width, height int) (*TUIDriver, error) {
	stdin, err := session.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdin pipe: %w", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdout pipe: %w", err)
	}

	stderr, err := session.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stderr pipe: %w", err)
	}

	driver := &TUIDriver{
		session:   session,
		stdin:     stdin,
		stdout:    stdout,
		stderr:    stderr,
		outputBuf: &bytes.Buffer{},
		width:     width,
		height:    height,
	}

	return driver, nil
}

// SendKey sends a single key to the TUI
func (d *TUIDriver) SendKey(key byte) error {
	_, err := d.stdin.Write([]byte{key})
	return err
}

// SendKeys sends multiple keys to the TUI
func (d *TUIDriver) SendKeys(keys string) error {
	_, err := d.stdin.Write([]byte(keys))
	return err
}

// SendEscape sends the Escape key
func (d *TUIDriver) SendEscape() error {
	return d.SendKey(0x1b) // ESC
}

// SendEnter sends the Enter key
func (d *TUIDriver) SendEnter() error {
	return d.SendKey('\n')
}

// SendCtrlC sends Ctrl+C (SIGINT)
func (d *TUIDriver) SendCtrlC() error {
	return d.SendKey(0x03) // ETX (Ctrl+C)
}

// SendArrowUp sends the up arrow key
func (d *TUIDriver) SendArrowUp() error {
	_, err := d.stdin.Write([]byte{0x1b, '[', 'A'})
	return err
}

// SendArrowDown sends the down arrow key
func (d *TUIDriver) SendArrowDown() error {
	_, err := d.stdin.Write([]byte{0x1b, '[', 'B'})
	return err
}

// SendArrowLeft sends the left arrow key
func (d *TUIDriver) SendArrowLeft() error {
	_, err := d.stdin.Write([]byte{0x1b, '[', 'D'})
	return err
}

// SendArrowRight sends the right arrow key
func (d *TUIDriver) SendArrowRight() error {
	_, err := d.stdin.Write([]byte{0x1b, '[', 'C'})
	return err
}

// SendTab sends the Tab key
func (d *TUIDriver) SendTab() error {
	return d.SendKey('\t')
}

// CaptureScreen reads the current screen output
func (d *TUIDriver) CaptureScreen(timeout time.Duration) (string, error) {
	// Clear previous buffer
	d.outputBuf.Reset()

	// Read with timeout
	done := make(chan error, 1)
	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := d.stdout.Read(buf)
			if n > 0 {
				d.outputBuf.Write(buf[:n])
			}
			if err != nil {
				done <- err
				return
			}
		}
	}()

	select {
	case <-time.After(timeout):
		return d.outputBuf.String(), nil
	case err := <-done:
		if err == io.EOF {
			return d.outputBuf.String(), nil
		}
		return d.outputBuf.String(), err
	}
}

// WaitForOutput waits for specific output to appear
func (d *TUIDriver) WaitForOutput(expected string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	d.outputBuf.Reset()

	buf := make([]byte, 1024)
	for time.Now().Before(deadline) {
		n, err := d.stdout.Read(buf)
		if n > 0 {
			d.outputBuf.Write(buf[:n])
			if bytes.Contains(d.outputBuf.Bytes(), []byte(expected)) {
				return nil
			}
		}
		if err != nil && err != io.EOF {
			// Continue on timeout
			time.Sleep(50 * time.Millisecond)
			continue
		}
		// Short sleep to prevent tight loop
		time.Sleep(50 * time.Millisecond)
	}

	return fmt.Errorf("timeout waiting for output containing %q. Got: %s", expected, d.outputBuf.String())
}

// Close closes the TUI driver and session
func (d *TUIDriver) Close() error {
	d.stdin.Close()
	return d.session.Close()
}

// Wait waits for the session to complete
func (d *TUIDriver) Wait() error {
	return d.session.Wait()
}

// Signal sends a signal to the session
func (d *TUIDriver) Signal(sig ssh.Signal) error {
	return d.session.Signal(sig)
}