go-llm/v2/sandbox/sandbox_test.go

package sandbox

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// --- Proxmox API mock server ---

type mockProxmoxHandler struct {
	containers map[int]ContainerStatus
	nextID     int
	tasks      map[string]string // taskID → exitstatus
}

func newMockProxmoxHandler() *mockProxmoxHandler {
	return &mockProxmoxHandler{
		containers: make(map[int]ContainerStatus),
		nextID:     200,
		tasks:      make(map[string]string),
	}
}

func (m *mockProxmoxHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path

	// Verify auth header is present.
	auth := r.Header.Get("Authorization")
	if !strings.HasPrefix(auth, "PVEAPIToken=") {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	w.Header().Set("Content-Type", "application/json")

	switch {
	case path == "/api2/json/cluster/nextid" && r.Method == http.MethodGet:
		m.handleNextID(w)

	case strings.HasSuffix(path, "/clone") && r.Method == http.MethodPost:
		m.handleClone(w, r)

	case strings.HasSuffix(path, "/config") && r.Method == http.MethodPut:
		m.handleConfig(w)

	case strings.HasSuffix(path, "/status/start") && r.Method == http.MethodPost:
		m.handleStart(w, r)

	case strings.HasSuffix(path, "/status/stop") && r.Method == http.MethodPost:
		m.handleStop(w, r)

	case strings.HasSuffix(path, "/status/current") && r.Method == http.MethodGet:
		m.handleStatusCurrent(w, r)

	case strings.HasSuffix(path, "/interfaces") && r.Method == http.MethodGet:
		m.handleInterfaces(w, r)

	case strings.Contains(path, "/tasks/") && strings.HasSuffix(path, "/status"):
		m.handleTaskStatus(w, r)

	case r.Method == http.MethodDelete && strings.Contains(path, "/lxc/"):
		m.handleDelete(w, r)

	case strings.HasSuffix(path, "/execute") && r.Method == http.MethodPost:
		m.handleExecute(w)

	default:
		http.Error(w, fmt.Sprintf("unhandled: %s %s", r.Method, path), http.StatusNotFound)
	}
}

func (m *mockProxmoxHandler) handleNextID(w http.ResponseWriter) {
	id := m.nextID
	m.nextID++
	json.NewEncoder(w).Encode(map[string]any{"data": id})
}

func (m *mockProxmoxHandler) handleClone(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	taskID := "UPID:pve:clone-task"
	m.tasks[taskID] = "OK"
	json.NewEncoder(w).Encode(map[string]any{"data": taskID})
}

func (m *mockProxmoxHandler) handleConfig(w http.ResponseWriter) {
	json.NewEncoder(w).Encode(map[string]any{"data": nil})
}

func (m *mockProxmoxHandler) handleStart(w http.ResponseWriter, r *http.Request) {
	// Extract VMID from path.
	parts := strings.Split(r.URL.Path, "/")
	for i, p := range parts {
		if p == "lxc" && i+1 < len(parts) {
			var vmid int
			fmt.Sscanf(parts[i+1], "%d", &vmid)
			m.containers[vmid] = ContainerStatus{Status: "running"}
			break
		}
	}
	taskID := "UPID:pve:start-task"
	m.tasks[taskID] = "OK"
	json.NewEncoder(w).Encode(map[string]any{"data": taskID})
}

func (m *mockProxmoxHandler) handleStop(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")
	for i, p := range parts {
		if p == "lxc" && i+1 < len(parts) {
			var vmid int
			fmt.Sscanf(parts[i+1], "%d", &vmid)
			m.containers[vmid] = ContainerStatus{Status: "stopped"}
			break
		}
	}
	taskID := "UPID:pve:stop-task"
	m.tasks[taskID] = "OK"
	json.NewEncoder(w).Encode(map[string]any{"data": taskID})
}

func (m *mockProxmoxHandler) handleStatusCurrent(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")
	for i, p := range parts {
		if p == "lxc" && i+1 < len(parts) {
			var vmid int
			fmt.Sscanf(parts[i+1], "%d", &vmid)
			status, ok := m.containers[vmid]
			if !ok {
				status = ContainerStatus{Status: "stopped"}
			}
			json.NewEncoder(w).Encode(map[string]any{"data": status})
			return
		}
	}
	http.Error(w, "not found", http.StatusNotFound)
}

func (m *mockProxmoxHandler) handleInterfaces(w http.ResponseWriter, r *http.Request) {
	ifaces := []map[string]string{
		{"name": "lo", "inet": "127.0.0.1/8"},
		{"name": "eth0", "inet": "10.99.1.5/16", "hwaddr": "AA:BB:CC:DD:EE:FF"},
	}
	json.NewEncoder(w).Encode(map[string]any{"data": ifaces})
}

func (m *mockProxmoxHandler) handleTaskStatus(w http.ResponseWriter, r *http.Request) {
	json.NewEncoder(w).Encode(map[string]any{
		"data": map[string]any{
			"status":     "stopped",
			"exitstatus": "OK",
		},
	})
}

func (m *mockProxmoxHandler) handleDelete(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")
	for i, p := range parts {
		if p == "lxc" && i+1 < len(parts) {
			var vmid int
			fmt.Sscanf(parts[i+1], "%d", &vmid)
			delete(m.containers, vmid)
			break
		}
	}
	taskID := "UPID:pve:delete-task"
	m.tasks[taskID] = "OK"
	json.NewEncoder(w).Encode(map[string]any{"data": taskID})
}

func (m *mockProxmoxHandler) handleExecute(w http.ResponseWriter) {
	json.NewEncoder(w).Encode(map[string]any{"data": ""})
}

// --- Test helpers ---

func newTestProxmoxClient(t *testing.T, handler *mockProxmoxHandler) (*ProxmoxClient, *httptest.Server) {
	t.Helper()
	server := httptest.NewTLSServer(handler)
	client := NewProxmoxClient(ProxmoxConfig{
		BaseURL:            server.URL,
		TokenID:            "test@pve!test-token",
		Secret:             "test-secret",
		Node:               "pve",
		TemplateID:         9000,
		Pool:               "sandbox-pool",
		Bridge:             "vmbr1",
		InsecureSkipVerify: true,
	})
	// Use the test server's TLS client.
	client.http = server.Client()
	return client, server
}

func generateTestSigner(t *testing.T) ssh.Signer {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate RSA key: %v", err)
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		t.Fatalf("create signer: %v", err)
	}
	return signer
}

// --- Proxmox client tests ---

func TestProxmoxNextAvailableID(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	id, err := client.NextAvailableID(context.Background())
	if err != nil {
		t.Fatalf("NextAvailableID: %v", err)
	}
	if id != 200 {
		t.Errorf("expected VMID 200, got %d", id)
	}

	// Second call should return 201.
	id2, err := client.NextAvailableID(context.Background())
	if err != nil {
		t.Fatalf("NextAvailableID (2nd): %v", err)
	}
	if id2 != 201 {
		t.Errorf("expected VMID 201, got %d", id2)
	}
}

func TestProxmoxCloneTemplate(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	err := client.CloneTemplate(context.Background(), 200, ContainerConfig{
		Hostname: "test-sandbox",
	})
	if err != nil {
		t.Fatalf("CloneTemplate: %v", err)
	}
}

func TestProxmoxContainerLifecycle(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	ctx := context.Background()

	// Start container.
	if err := client.StartContainer(ctx, 200); err != nil {
		t.Fatalf("StartContainer: %v", err)
	}

	// Get status — should be running.
	status, err := client.GetContainerStatus(ctx, 200)
	if err != nil {
		t.Fatalf("GetContainerStatus: %v", err)
	}
	if status.Status != "running" {
		t.Errorf("expected status 'running', got %q", status.Status)
	}

	// Stop container.
	if err := client.StopContainer(ctx, 200); err != nil {
		t.Fatalf("StopContainer: %v", err)
	}

	// Get status — should be stopped.
	status, err = client.GetContainerStatus(ctx, 200)
	if err != nil {
		t.Fatalf("GetContainerStatus: %v", err)
	}
	if status.Status != "stopped" {
		t.Errorf("expected status 'stopped', got %q", status.Status)
	}
}

func TestProxmoxGetContainerIP(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	ip, err := client.GetContainerIP(context.Background(), 200)
	if err != nil {
		t.Fatalf("GetContainerIP: %v", err)
	}
	if ip != "10.99.1.5" {
		t.Errorf("expected IP 10.99.1.5, got %q", ip)
	}
}

func TestProxmoxDestroyContainer(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	ctx := context.Background()

	// Start it first so it has a status.
	if err := client.StartContainer(ctx, 200); err != nil {
		t.Fatalf("StartContainer: %v", err)
	}

	// Destroy it.
	if err := client.DestroyContainer(ctx, 200); err != nil {
		t.Fatalf("DestroyContainer: %v", err)
	}

	// Container should be gone from the handler's map.
	if _, exists := handler.containers[200]; exists {
		t.Error("container 200 should have been deleted")
	}
}

func TestProxmoxConfigureContainer(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	err := client.ConfigureContainer(context.Background(), 200, ContainerConfig{
		CPUs:     2,
		MemoryMB: 2048,
	})
	if err != nil {
		t.Fatalf("ConfigureContainer: %v", err)
	}
}

func TestProxmoxEnableDisableInternet(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	ctx := context.Background()

	if err := client.EnableInternet(ctx, "10.99.1.5"); err != nil {
		t.Fatalf("EnableInternet: %v", err)
	}

	if err := client.DisableInternet(ctx, "10.99.1.5"); err != nil {
		t.Fatalf("DisableInternet: %v", err)
	}
}

func TestProxmoxAuthRequired(t *testing.T) {
	// Mock that rejects requests without a valid token.
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		if auth != "PVEAPIToken=valid@pve!tok=secret123" {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{"data": 200})
	})
	server := httptest.NewTLSServer(handler)
	defer server.Close()

	// Client with wrong credentials should fail.
	client := NewProxmoxClient(ProxmoxConfig{
		BaseURL:            server.URL,
		TokenID:            "wrong@pve!tok",
		Secret:             "wrong",
		Node:               "pve",
		InsecureSkipVerify: true,
	})
	client.http = server.Client()

	_, err := client.NextAvailableID(context.Background())
	if err == nil {
		t.Fatal("expected error with wrong auth, got nil")
	}
	if !strings.Contains(err.Error(), "401") {
		t.Errorf("expected 401 error, got: %v", err)
	}

	// Client with correct credentials should succeed.
	client2 := NewProxmoxClient(ProxmoxConfig{
		BaseURL:            server.URL,
		TokenID:            "valid@pve!tok",
		Secret:             "secret123",
		Node:               "pve",
		InsecureSkipVerify: true,
	})
	client2.http = server.Client()

	id, err := client2.NextAvailableID(context.Background())
	if err != nil {
		t.Fatalf("expected success with correct auth, got: %v", err)
	}
	if id != 200 {
		t.Errorf("expected VMID 200, got %d", id)
	}
}

func TestProxmoxContextCancellation(t *testing.T) {
	handler := newMockProxmoxHandler()
	client, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // Cancel immediately.

	_, err := client.NextAvailableID(ctx)
	if err == nil {
		t.Fatal("expected error with cancelled context, got nil")
	}
}

// --- SSH executor tests (unit tests without real SSH) ---

func TestSSHExecutorDefaults(t *testing.T) {
	signer := generateTestSigner(t)
	exec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	if exec.config.User != "sandbox" {
		t.Errorf("expected default user 'sandbox', got %q", exec.config.User)
	}
	if exec.config.ConnectTimeout != 10e9 {
		t.Errorf("expected default connect timeout 10s, got %v", exec.config.ConnectTimeout)
	}
	if exec.config.CommandTimeout != 60e9 {
		t.Errorf("expected default command timeout 60s, got %v", exec.config.CommandTimeout)
	}
}

func TestSSHExecutorNotConnected(t *testing.T) {
	signer := generateTestSigner(t)
	exec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	_, err := exec.Exec(context.Background(), "echo hello")
	if err == nil {
		t.Fatal("expected error when not connected, got nil")
	}
	if !strings.Contains(err.Error(), "not connected") {
		t.Errorf("expected 'not connected' error, got: %v", err)
	}
}

func TestSSHExecutorUploadNotConnected(t *testing.T) {
	signer := generateTestSigner(t)
	exec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	err := exec.Upload(context.Background(), strings.NewReader("test"), "/tmp/test", 0644)
	if err == nil {
		t.Fatal("expected error when not connected, got nil")
	}
}

func TestSSHExecutorDownloadNotConnected(t *testing.T) {
	signer := generateTestSigner(t)
	exec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	_, err := exec.Download(context.Background(), "/tmp/test")
	if err == nil {
		t.Fatal("expected error when not connected, got nil")
	}
}

func TestSSHExecutorIsConnected(t *testing.T) {
	signer := generateTestSigner(t)
	exec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	if exec.IsConnected() {
		t.Error("should not be connected initially")
	}
}

func TestSSHExecutorCloseIdempotent(t *testing.T) {
	signer := generateTestSigner(t)
	exec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	// Close without connecting should not error.
	if err := exec.Close(); err != nil {
		t.Errorf("Close on unconnected executor: %v", err)
	}
}

// --- LoadSSHKey / ParseSSHKey tests ---

func TestLoadSSHKeyNotFound(t *testing.T) {
	_, err := LoadSSHKey("/nonexistent/path/to/key")
	if err == nil {
		t.Fatal("expected error for nonexistent key, got nil")
	}
}

func TestParseSSHKeyInvalid(t *testing.T) {
	_, err := ParseSSHKey([]byte("not a valid PEM key"))
	if err == nil {
		t.Fatal("expected error for invalid key, got nil")
	}
}

// --- Sandbox / Manager tests (using mock Proxmox, no real SSH) ---

func TestManagerRequiresSigner(t *testing.T) {
	_, err := NewManager(Config{})
	if err == nil {
		t.Fatal("expected error when no SSH signer provided")
	}
	if !strings.Contains(err.Error(), "SSH signer") {
		t.Errorf("expected SSH signer error, got: %v", err)
	}
}

func TestSandboxDestroyClosesConnections(t *testing.T) {
	handler := newMockProxmoxHandler()
	_, server := newTestProxmoxClient(t, handler)
	defer server.Close()

	signer := generateTestSigner(t)

	proxmoxClient := NewProxmoxClient(ProxmoxConfig{
		BaseURL:            server.URL,
		TokenID:            "test@pve!test-token",
		Secret:             "test-secret",
		Node:               "pve",
		TemplateID:         9000,
		InsecureSkipVerify: true,
	})
	proxmoxClient.http = server.Client()

	// Start a container so destroy can check its status.
	ctx := context.Background()
	if err := proxmoxClient.StartContainer(ctx, 200); err != nil {
		t.Fatalf("start: %v", err)
	}

	sshExec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	sb := &Sandbox{
		ID:       200,
		IP:       "10.99.1.5",
		Internet: false,
		proxmox:  proxmoxClient,
		ssh:      sshExec,
	}

	// Destroy should succeed even with unconnected SSH (no SFTP/SSH to close).
	if err := sb.Destroy(ctx); err != nil {
		t.Fatalf("Destroy: %v", err)
	}
}

func TestSandboxWriteFileAndReadFileRequireConnection(t *testing.T) {
	signer := generateTestSigner(t)
	sshExec := NewSSHExecutor("10.99.1.5", SSHConfig{Signer: signer})

	sb := &Sandbox{ssh: sshExec}

	err := sb.WriteFile(context.Background(), "/tmp/test.txt", "hello")
	if err == nil {
		t.Fatal("expected error when SSH not connected")
	}

	_, err = sb.ReadFile(context.Background(), "/tmp/test.txt")
	if err == nil {
		t.Fatal("expected error when SSH not connected")
	}
}

func TestContainerConfigDefaults(t *testing.T) {
	// Verify that zero-value createOpts get proper defaults in the Create flow.
	o := &createOpts{}

	if o.cpus != 0 {
		t.Errorf("expected zero cpus, got %d", o.cpus)
	}

	// Apply options.
	WithCPUs(2)(o)
	WithMemoryMB(2048)(o)
	WithDiskGB(16)(o)
	WithHostname("test")(o)
	WithInternet(true)(o)

	if o.cpus != 2 {
		t.Errorf("expected cpus=2, got %d", o.cpus)
	}
	if o.memoryMB != 2048 {
		t.Errorf("expected memoryMB=2048, got %d", o.memoryMB)
	}
	if o.diskGB != 16 {
		t.Errorf("expected diskGB=16, got %d", o.diskGB)
	}
	if o.hostname != "test" {
		t.Errorf("expected hostname='test', got %q", o.hostname)
	}
	if !o.internet {
		t.Error("expected internet=true")
	}
}

func TestExecResultFields(t *testing.T) {
	r := ExecResult{Output: "hello\n", ExitCode: 0}
	if r.Output != "hello\n" {
		t.Errorf("unexpected output: %q", r.Output)
	}
	if r.ExitCode != 0 {
		t.Errorf("unexpected exit code: %d", r.ExitCode)
	}
}