codespace/internal/shell/manager_test.go

package shell

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

import (
	"codespace/internal/util"
)

func startTestShell(t *testing.T) (*LocalManager, string, io.WriteCloser, Subscription) {
	t.Helper()
	mgr := NewManager("bash", []string{"-i"})
	root := t.TempDir()
	shellID, err := mgr.Start("test-ws", root)
	if err != nil {
		t.Fatalf("Start failed: %v", err)
	}

	stdin, err := mgr.Stdin("test-ws", shellID)
	if err != nil {
		t.Fatalf("Stdin failed: %v", err)
	}

	sub, err := mgr.Subscribe("test-ws", shellID)
	if err != nil {
		t.Fatalf("Subscribe failed: %v", err)
	}

	t.Cleanup(func() {
		sub.Close()
		_ = mgr.Stop("test-ws", shellID)
	})

	return mgr, shellID, stdin, sub
}

func waitForMarker(t *testing.T, sub Subscription, stdin io.Writer, input, marker string) string {
	t.Helper()

	if _, err := stdin.Write([]byte(input)); err != nil {
		t.Fatalf("failed to write to stdin: %v", err)
	}

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

	var out bytes.Buffer
	for {
		select {
		case chunk, ok := <-sub.Output():
			if !ok {
				t.Fatalf("subscription closed before seeing %q; output so far:\n%s", marker, out.String())
			}
			out.Write(chunk)
			if strings.Contains(out.String(), marker) {
				return out.String()
			}
		case <-ctx.Done():
			t.Fatalf("timed out waiting for %q; output so far:\n%s", marker, out.String())
		}
	}
}

func TestShellPTYIsRealTerminal(t *testing.T) {
	mgr, shellID, stdin, sub := startTestShell(t)
	_ = mgr
	_ = shellID

	out := waitForMarker(t, sub, stdin, "echo HELLO_PTY_TEST\n", "HELLO_PTY_TEST")

	if !strings.Contains(out, "echo HELLO_PTY_TEST") {
		t.Errorf("expected command echo in PTY output, got:\n%s", out)
	}
	if !strings.Contains(out, "HELLO_PTY_TEST") {
		t.Errorf("expected command output in PTY output, got:\n%s", out)
	}
}

func TestShellPTYTermEnv(t *testing.T) {
	_, _, stdin, sub := startTestShell(t)

	out := waitForMarker(t, sub, stdin, "echo \"$TERM\"\n", "xterm-256color")
	if !strings.Contains(out, "xterm-256color") {
		t.Errorf("expected TERM=xterm-256color in output, got:\n%s", out)
	}
}

func TestShellResizeUpdatesPTYSize(t *testing.T) {
	mgr, shellID, stdin, sub := startTestShell(t)
	_ = sub

	if err := mgr.Resize("test-ws", shellID, 120, 40); err != nil {
		t.Fatalf("Resize failed: %v", err)
	}

	out := waitForMarker(t, sub, stdin, "stty size\n", "40 120")
	if !strings.Contains(out, "40 120") {
		t.Errorf("expected stty size to report 40 120, got:\n%s", out)
	}
}

func TestShellResizeRejectsInvalidSize(t *testing.T) {
	mgr, shellID, _, _ := startTestShell(t)

	if err := mgr.Resize("test-ws", shellID, 0, 40); err == nil {
		t.Fatal("expected error for invalid size")
	} else if util.CodeOf(err) != util.CodeBadRequest {
		t.Errorf("expected CodeBadRequest, got %v", util.CodeOf(err))
	}
}

func TestShellResizeRejectsMissingSession(t *testing.T) {
	mgr := NewManager("bash", []string{"-i"})

	if err := mgr.Resize("missing-ws", "missing-shell", 80, 24); err == nil {
		t.Fatal("expected error for missing session")
	} else if util.CodeOf(err) != util.CodeNotFound {
		t.Errorf("expected CodeNotFound, got %v", util.CodeOf(err))
	}
}

func TestShellMultiInstance(t *testing.T) {
	mgr := NewManager("bash", []string{"-i"})
	root := t.TempDir()

	shellID1, err := mgr.Start("test-ws", root)
	if err != nil {
		t.Fatalf("Start shell1 failed: %v", err)
	}
	shellID2, err := mgr.Start("test-ws", root)
	if err != nil {
		t.Fatalf("Start shell2 failed: %v", err)
	}

	if shellID1 == shellID2 {
		t.Fatalf("expected distinct shell IDs, got %q and %q", shellID1, shellID2)
	}

	stdin1, err := mgr.Stdin("test-ws", shellID1)
	if err != nil {
		t.Fatalf("Stdin shell1 failed: %v", err)
	}
	stdin2, err := mgr.Stdin("test-ws", shellID2)
	if err != nil {
		t.Fatalf("Stdin shell2 failed: %v", err)
	}

	sub1, err := mgr.Subscribe("test-ws", shellID1)
	if err != nil {
		t.Fatalf("Subscribe shell1 failed: %v", err)
	}
	sub2, err := mgr.Subscribe("test-ws", shellID2)
	if err != nil {
		t.Fatalf("Subscribe shell2 failed: %v", err)
	}

	t.Cleanup(func() {
		sub1.Close()
		sub2.Close()
		_ = mgr.Stop("test-ws", shellID1)
		_ = mgr.Stop("test-ws", shellID2)
	})

	list := mgr.List("test-ws")
	if len(list) != 2 {
		t.Fatalf("List returned %d shells, want 2", len(list))
	}

	out1 := waitForMarker(t, sub1, stdin1, "echo SHELL_ONE\n", "SHELL_ONE")
	if strings.Contains(out1, "SHELL_TWO") {
		t.Errorf("shell1 output leaked shell2 output: %s", out1)
	}

	out2 := waitForMarker(t, sub2, stdin2, "echo SHELL_TWO\n", "SHELL_TWO")
	if strings.Contains(out2, "SHELL_ONE") {
		t.Errorf("shell2 output leaked shell1 output: %s", out2)
	}

	if err := mgr.Stop("test-ws", shellID1); err != nil {
		t.Fatalf("Stop shell1 failed: %v", err)
	}

	if _, err := mgr.Stdin("test-ws", shellID1); err == nil {
		t.Fatal("expected shell1 stdin to be unavailable after stop")
	}

	// shell2 should still work.
	_ = waitForMarker(t, sub2, stdin2, "echo STILL_ALIVE\n", "STILL_ALIVE")

	list = mgr.List("test-ws")
	if len(list) != 1 {
		t.Fatalf("List returned %d shells after stop, want 1", len(list))
	}
	if list[0].ShellID != shellID2 {
		t.Fatalf("List returned shell %q, want %q", list[0].ShellID, shellID2)
	}
}