oikos/internal/mcp/sshexec_test.go

package mcp

// Streaming tests for sshExecStream against a real SSH endpoint. Guarded by
// OIKOS_SSH_TEST_HOST — skipped when unset. Run with:
//
//	OIKOS_SSH_TEST_HOST=localhost OIKOS_SSH_USER=$USER \
//	OIKOS_SSH_KEY_PATH=~/.ssh/id_ed25519 go test ./internal/mcp/ -run TestSSHExecStream
//
// These matter because the whole point of the change is behaviour that only
// appears over time: that output arrives *before* the command exits, and that
// a command killed mid-flight still leaves what it printed.

import (
	"context"
	"os"
	"strings"
	"sync"
	"testing"
	"time"
)

func sshTestHost(t *testing.T) string {
	t.Helper()
	host := os.Getenv("OIKOS_SSH_TEST_HOST")
	if host == "" {
		t.Skip("OIKOS_SSH_TEST_HOST not set — skipping live SSH test")
	}
	return host
}

// The core claim: chunks reach the sink while the command is still running,
// not in one lump at the end. A command that prints, sleeps, then prints must
// deliver its first chunk well before it exits.
func TestSSHExecStreamDeliversOutputBeforeExit(t *testing.T) {
	host := sshTestHost(t)

	var (
		mu     sync.Mutex
		chunks []string
		firstA time.Time
	)
	sink := func(stream string, chunk []byte) {
		mu.Lock()
		defer mu.Unlock()
		if firstA.IsZero() {
			firstA = time.Now()
		}
		chunks = append(chunks, string(chunk))
	}

	start := time.Now()
	out, err := sshExecStream(context.Background(), host, os.Getenv("OIKOS_SSH_USER"),
		"echo FIRST; sleep 2; echo SECOND", sink)
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("sshExecStream: %v (out=%q)", err, out)
	}

	mu.Lock()
	joined := strings.Join(chunks, "")
	firstAt := firstA.Sub(start)
	mu.Unlock()

	if !strings.Contains(out, "FIRST") || !strings.Contains(out, "SECOND") {
		t.Errorf("combined output lost content: %q", out)
	}
	if !strings.Contains(joined, "FIRST") || !strings.Contains(joined, "SECOND") {
		t.Errorf("sink did not receive the full output: %q", joined)
	}
	if elapsed < 2*time.Second {
		t.Fatalf("command returned in %v — the sleep did not run, test is not measuring what it claims", elapsed)
	}
	// The first chunk must land near the start, not at the end.
	if firstAt > elapsed/2 {
		t.Errorf("first chunk arrived after %v of a %v command — output is still being buffered to the end",
			firstAt, elapsed)
	}
}

// stderr must reach the sink too, and land in the combined output, matching
// what CombinedOutput used to return.
func TestSSHExecStreamCapturesBothStreams(t *testing.T) {
	host := sshTestHost(t)

	var mu sync.Mutex
	seen := map[string]bool{}
	sink := func(stream string, chunk []byte) {
		mu.Lock()
		defer mu.Unlock()
		seen[stream] = true
	}

	out, err := sshExecStream(context.Background(), host, os.Getenv("OIKOS_SSH_USER"),
		"echo TO_STDOUT; echo TO_STDERR 1>&2", sink)
	if err != nil {
		t.Fatalf("sshExecStream: %v (out=%q)", err, out)
	}

	if !strings.Contains(out, "TO_STDOUT") || !strings.Contains(out, "TO_STDERR") {
		t.Errorf("combined output missing a stream: %q", out)
	}
	mu.Lock()
	defer mu.Unlock()
	if !seen["stdout"] {
		t.Error("sink never saw a stdout chunk")
	}
	if !seen["stderr"] {
		t.Error("sink never saw a stderr chunk")
	}
}

// A cancelled command used to return "" — everything it had printed was
// thrown away. The hung case is exactly when that output is worth having.
func TestSSHExecStreamKeepsPartialOutputOnCancel(t *testing.T) {
	host := sshTestHost(t)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	out, err := sshExecStream(ctx, host, os.Getenv("OIKOS_SSH_USER"),
		"echo BEFORE_HANG; sleep 30; echo NEVER", nil)

	if err == nil {
		t.Fatal("expected a context error for a command that outlives the deadline")
	}
	if !strings.Contains(out, "BEFORE_HANG") {
		t.Errorf("partial output was discarded on cancel: %q", out)
	}
	if strings.Contains(out, "NEVER") {
		t.Errorf("command should not have completed: %q", out)
	}
}

// A nil sink must behave exactly as the old CombinedOutput path did.
func TestSSHExecNilSinkStillReturnsOutput(t *testing.T) {
	host := sshTestHost(t)

	out, err := sshExec(context.Background(), host, os.Getenv("OIKOS_SSH_USER"), "echo PLAIN")
	if err != nil {
		t.Fatalf("sshExec: %v", err)
	}
	if out != "PLAIN" {
		t.Errorf("out = %q, want %q (output is trimmed)", out, "PLAIN")
	}
}

// A non-zero exit must surface as an error while still returning the output.
func TestSSHExecStreamNonZeroExitIsAnError(t *testing.T) {
	host := sshTestHost(t)

	out, err := sshExecStream(context.Background(), host, os.Getenv("OIKOS_SSH_USER"),
		"echo PRINTED_THEN_FAILED; exit 3", nil)
	if err == nil {
		t.Fatal("a non-zero exit that printed output must still be an error")
	}
	if !strings.Contains(out, "PRINTED_THEN_FAILED") {
		t.Errorf("output lost on failure: %q", out)
	}
}