core/server/sandbox/workspace_module_test.go

package sandbox

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"go.starlark.net/starlark"
)

// ── Test helpers ──────────────────────────────

func execWithWorkspace(t *testing.T, script string) (*Result, string, error) {
	t.Helper()
	tmpDir := t.TempDir()
	ctx := context.Background()
	mod := BuildWorkspaceModule(ctx, WorkspaceModuleConfig{
		PackageID:     "test-pkg",
		WorkspaceRoot: tmpDir,
	})
	sb := New(DefaultConfig())
	result, err := sb.Exec(ctx, "test.star", script, map[string]starlark.Value{
		"workspace": mod,
	})
	return result, tmpDir, err
}

func execWithWorkspaceQuota(t *testing.T, quotaMB int, script string) (*Result, string, error) {
	t.Helper()
	tmpDir := t.TempDir()
	ctx := context.Background()
	mod := BuildWorkspaceModule(ctx, WorkspaceModuleConfig{
		PackageID:     "test-pkg",
		WorkspaceRoot: tmpDir,
		QuotaMB:       quotaMB,
	})
	sb := New(DefaultConfig())
	result, err := sb.Exec(ctx, "test.star", script, map[string]starlark.Value{
		"workspace": mod,
	})
	return result, tmpDir, err
}

// ── Create ────────────────────────────────────

func TestWorkspaceCreate_ReturnsPath(t *testing.T) {
	result, tmpDir, err := execWithWorkspace(t, `
path = workspace.create("myrepo")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	got := string(result.Globals["path"].(starlark.String))
	expected := filepath.Join(tmpDir, "test-pkg", "myrepo")
	if got != expected {
		t.Errorf("path = %q, want %q", got, expected)
	}
}

func TestWorkspaceCreate_Idempotent(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
path1 = workspace.create("myrepo")
path2 = workspace.create("myrepo")
same = (path1 == path2)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	if result.Globals["same"] != starlark.True {
		t.Error("expected same path on second create")
	}
}

func TestWorkspaceCreate_DirectoryExists(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
path = workspace.create("build_cache")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	dir := string(result.Globals["path"].(starlark.String))
	info, statErr := os.Stat(dir)
	if statErr != nil {
		t.Fatalf("directory does not exist: %v", statErr)
	}
	if !info.IsDir() {
		t.Error("expected a directory")
	}
}

// ── Path ──────────────────────────────────────

func TestWorkspacePath_Exists(t *testing.T) {
	result, tmpDir, err := execWithWorkspace(t, `
workspace.create("myrepo")
path = workspace.path("myrepo")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	got := string(result.Globals["path"].(starlark.String))
	expected := filepath.Join(tmpDir, "test-pkg", "myrepo")
	if got != expected {
		t.Errorf("path = %q, want %q", got, expected)
	}
}

func TestWorkspacePath_NotFound(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
path = workspace.path("nonexistent")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	if result.Globals["path"] != starlark.None {
		t.Errorf("expected None, got %v", result.Globals["path"])
	}
}

// ── List ──────────────────────────────────────

func TestWorkspaceList_Empty(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
names = workspace.list()
count = len(names)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	count, _ := starlark.AsInt32(result.Globals["count"])
	if count != 0 {
		t.Errorf("count = %d, want 0", count)
	}
}

func TestWorkspaceList_Multiple(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
workspace.create("alpha")
workspace.create("beta")
workspace.create("gamma")
names = workspace.list()
count = len(names)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	count, _ := starlark.AsInt32(result.Globals["count"])
	if count != 3 {
		t.Errorf("count = %d, want 3", count)
	}
}

// ── Delete ────────────────────────────────────

func TestWorkspaceDelete_Removes(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
workspace.create("temp")
workspace.delete("temp")
path = workspace.path("temp")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	if result.Globals["path"] != starlark.None {
		t.Error("expected None after delete")
	}
}

func TestWorkspaceDelete_Idempotent(t *testing.T) {
	_, _, err := execWithWorkspace(t, `
workspace.delete("nonexistent")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
}

// ── Usage ─────────────────────────────────────

func TestWorkspaceUsage_EmptyDir(t *testing.T) {
	result, _, err := execWithWorkspace(t, `
workspace.create("empty")
bytes = workspace.usage("empty")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	usage, _ := starlark.AsInt32(result.Globals["bytes"])
	if usage != 0 {
		t.Errorf("usage = %d, want 0", usage)
	}
}

func TestWorkspaceUsage_WithFiles(t *testing.T) {
	tmpDir := t.TempDir()
	ctx := context.Background()

	// Pre-create workspace directory with a file.
	wsDir := filepath.Join(tmpDir, "test-pkg", "data")
	if err := os.MkdirAll(wsDir, 0755); err != nil {
		t.Fatal(err)
	}
	// Write a 1024-byte file.
	if err := os.WriteFile(filepath.Join(wsDir, "file.bin"), make([]byte, 1024), 0644); err != nil {
		t.Fatal(err)
	}

	mod := BuildWorkspaceModule(ctx, WorkspaceModuleConfig{
		PackageID:     "test-pkg",
		WorkspaceRoot: tmpDir,
	})
	sb := New(DefaultConfig())
	result, err := sb.Exec(ctx, "test.star", `bytes = workspace.usage("data")`, map[string]starlark.Value{
		"workspace": mod,
	})
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	usage, _ := starlark.AsInt32(result.Globals["bytes"])
	if usage != 1024 {
		t.Errorf("usage = %d, want 1024", usage)
	}
}

func TestWorkspaceUsage_NonExistent(t *testing.T) {
	_, _, err := execWithWorkspace(t, `
bytes = workspace.usage("nope")
`)
	if err == nil {
		t.Fatal("expected error for nonexistent workspace")
	}
	if !strings.Contains(err.Error(), "does not exist") {
		t.Errorf("unexpected error: %v", err)
	}
}

// ── Name Validation ──────────────────────────

func TestWorkspaceNameValidation_BadNames(t *testing.T) {
	badNames := []struct {
		name string
		desc string
	}{
		{"../escape", "traversal"},
		{"UPPER", "uppercase"},
		{"", "empty"},
		{"has space", "space"},
		{"has/slash", "slash"},
		{".dotstart", "dot prefix"},
		{"0numstart", "digit prefix"},
		{strings.Repeat("a", 64), "too long"},
	}

	for _, tc := range badNames {
		t.Run(tc.desc, func(t *testing.T) {
			script := fmt.Sprintf(`workspace.path(%q)`, tc.name)
			_, _, err := execWithWorkspace(t, script)
			if err == nil {
				t.Errorf("expected error for name %q", tc.name)
			}
		})
	}
}

func TestWorkspaceNameValidation_GoodNames(t *testing.T) {
	goodNames := []string{"myrepo", "build_cache", "a", "a0_b1_c2"}

	for _, name := range goodNames {
		t.Run(name, func(t *testing.T) {
			script := fmt.Sprintf(`path = workspace.path(%q)`, name)
			_, _, err := execWithWorkspace(t, script)
			if err != nil {
				t.Errorf("unexpected error for name %q: %v", name, err)
			}
		})
	}
}

// ── Quota ─────────────────────────────────────

func TestWorkspaceQuota_Enforced(t *testing.T) {
	tmpDir := t.TempDir()
	ctx := context.Background()

	// Pre-fill the package directory with >1 MB of data.
	pkgDir := filepath.Join(tmpDir, "test-pkg", "existing")
	if err := os.MkdirAll(pkgDir, 0755); err != nil {
		t.Fatal(err)
	}
	bigFile := make([]byte, 1100*1024) // ~1.07 MB
	if err := os.WriteFile(filepath.Join(pkgDir, "big.bin"), bigFile, 0644); err != nil {
		t.Fatal(err)
	}

	mod := BuildWorkspaceModule(ctx, WorkspaceModuleConfig{
		PackageID:     "test-pkg",
		WorkspaceRoot: tmpDir,
		QuotaMB:       1,
	})
	sb := New(DefaultConfig())
	_, err := sb.Exec(ctx, "test.star", `workspace.create("newone")`, map[string]starlark.Value{
		"workspace": mod,
	})
	if err == nil {
		t.Fatal("expected quota exceeded error")
	}
	if !strings.Contains(err.Error(), "quota exceeded") {
		t.Errorf("unexpected error: %v", err)
	}
}

func TestWorkspaceQuota_Unlimited(t *testing.T) {
	tmpDir := t.TempDir()
	ctx := context.Background()

	// Pre-fill with data — should not matter with quota=0.
	pkgDir := filepath.Join(tmpDir, "test-pkg", "existing")
	if err := os.MkdirAll(pkgDir, 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(pkgDir, "big.bin"), make([]byte, 2*1024*1024), 0644); err != nil {
		t.Fatal(err)
	}

	mod := BuildWorkspaceModule(ctx, WorkspaceModuleConfig{
		PackageID:     "test-pkg",
		WorkspaceRoot: tmpDir,
		QuotaMB:       0, // unlimited
	})
	sb := New(DefaultConfig())
	_, err := sb.Exec(ctx, "test.star", `workspace.create("newone")`, map[string]starlark.Value{
		"workspace": mod,
	})
	if err != nil {
		t.Fatalf("expected success with unlimited quota: %v", err)
	}
}