TTY

internal/terminal/session_test.go

@@ -0,0 +1,178 @@
+package terminal
+
+import (
+	"context"
+	"errors"
+	"runtime"
+	"strings"
+	"testing"
+	"time"
+)
+
+// testShell 返回各平台一个可交互读 stdin 的 shell。
+func testShell() string {
+	if runtime.GOOS == "windows" {
+		return "cmd"
+	}
+	return "/bin/sh"
+}
+
+func newTestSup(t *testing.T, max int) *Supervisor {
+	t.Helper()
+	sup := NewSupervisor(SupervisorConfig{
+		Shell:       testShell(),
+		MaxSessions: max,
+		KillGrace:   2 * time.Second,
+		EnvWhitelist: []string{
+			"PATH", "HOME", "TERM",
+			"SystemRoot", "TEMP", "TMP", "USERPROFILE", "ComSpec", "PATHEXT",
+		},
+	}, nil)
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		sup.CloseAll(ctx)
+	})
+	return sup
+}
+
+// readUntil 在独立 goroutine 中持续读会话输出，直到累计内容包含 marker 或超时。
+// 阻塞的 Read goroutine 会在会话被关闭（defer/Cleanup）后随 Read 报错而退出。
+func readUntil(sess *Session, marker string, d time.Duration) string {
+	out := make(chan []byte, 16)
+	go func() {
+		buf := make([]byte, 4096)
+		for {
+			n, err := sess.Read(buf)
+			if n > 0 {
+				b := make([]byte, n)
+				copy(b, buf[:n])
+				out <- b
+			}
+			if err != nil {
+				close(out)
+				return
+			}
+		}
+	}()
+	deadline := time.After(d)
+	var acc strings.Builder
+	for {
+		select {
+		case b, ok := <-out:
+			if !ok {
+				return acc.String()
+			}
+			acc.Write(b)
+			if strings.Contains(acc.String(), marker) {
+				return acc.String()
+			}
+		case <-deadline:
+			return acc.String()
+		}
+	}
+}
+
+// TestOpen_ReadsShellOutput 验证开会话后向 shell 写命令能读回其输出。
+// Unix 走真 PTY；Windows 走降级管道实现。
+func TestOpen_ReadsShellOutput(t *testing.T) {
+	sup := newTestSup(t, 5)
+	sess, err := sup.Open(OpenParams{})
+	if err != nil {
+		t.Fatalf("open: %v", err)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	defer sup.Close(ctx, sess.ID)
+
+	const marker = "ACWMARK42"
+	cmd := "echo " + marker + "\n"
+	if runtime.GOOS == "windows" {
+		cmd = "echo " + marker + "\r\n"
+	}
+	if _, err := sess.Write([]byte(cmd)); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+
+	got := readUntil(sess, marker, 5*time.Second)
+	if !strings.Contains(got, marker) {
+		t.Fatalf("marker %q not found in shell output; got %q", marker, got)
+	}
+}
+
+// TestOpen_MaxSessions 验证达到上限后再开会话返回 ErrTooManySessions。
+func TestOpen_MaxSessions(t *testing.T) {
+	sup := newTestSup(t, 1)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	s1, err := sup.Open(OpenParams{})
+	if err != nil {
+		t.Fatalf("first open: %v", err)
+	}
+	defer sup.Close(ctx, s1.ID)
+
+	if _, err := sup.Open(OpenParams{}); !errors.Is(err, ErrTooManySessions) {
+		t.Fatalf("expected ErrTooManySessions, got %v", err)
+	}
+	if got := sup.Count(); got != 1 {
+		t.Fatalf("expected 1 active session, got %d", got)
+	}
+}
+
+// TestClose_TerminatesSession 验证 Close 终止 shell（done 关闭）并从映射移除。
+func TestClose_TerminatesSession(t *testing.T) {
+	sup := newTestSup(t, 5)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	sess, err := sup.Open(OpenParams{})
+	if err != nil {
+		t.Fatalf("open: %v", err)
+	}
+	if got := sup.Count(); got != 1 {
+		t.Fatalf("expected 1 active session, got %d", got)
+	}
+
+	sup.Close(ctx, sess.ID)
+
+	select {
+	case <-sess.Done():
+	case <-time.After(8 * time.Second):
+		t.Fatal("session not terminated within 8s after Close")
+	}
+
+	// monitor 在 done 关闭后异步移除映射，轮询直至归零。
+	deadline := time.Now().Add(3 * time.Second)
+	for sup.Count() != 0 {
+		if time.Now().After(deadline) {
+			t.Fatalf("session not removed from supervisor; count=%d", sup.Count())
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+}
+
+// TestCloseAll_ClosesEverySession 验证 CloseAll 终止全部会话。
+func TestCloseAll_ClosesEverySession(t *testing.T) {
+	sup := newTestSup(t, 5)
+	for i := 0; i < 3; i++ {
+		if _, err := sup.Open(OpenParams{}); err != nil {
+			t.Fatalf("open %d: %v", i, err)
+		}
+	}
+	if got := sup.Count(); got != 3 {
+		t.Fatalf("expected 3 active sessions, got %d", got)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
+	defer cancel()
+	sup.CloseAll(ctx)
+
+	deadline := time.Now().Add(3 * time.Second)
+	for sup.Count() != 0 {
+		if time.Now().After(deadline) {
+			t.Fatalf("sessions not all closed; count=%d", sup.Count())
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+}