package workspace_test import ( "context" "net/url" "os" "strconv" "strings" "testing" "time" "github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/volume" "github.com/yaoapp/yao/workspace" ) type poolConfig struct { Name string Addr string } func testPools() []poolConfig { pools := []poolConfig{ {Name: "local", Addr: "local"}, } if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" { name := taiIDFromAddr(addr) pools = append(pools, poolConfig{Name: name, Addr: addr}) } return pools } func taiIDFromAddr(addr string) string { addr = strings.TrimSpace(addr) if addr == "local" || addr == "" { return "local" } if !strings.Contains(addr, "://") { addr = "tai://" + addr } u, err := url.Parse(addr) if err != nil { return addr } h := u.Hostname() if h == "" { return addr } if p := u.Port(); p != "" { return h + "-" + p } return h } func ensureRegistry(tb testing.TB) { tb.Helper() registry.Init(nil) } func setupManagerForPool(tb testing.TB, pc poolConfig) *workspace.Manager { tb.Helper() ensureRegistry(tb) registerForTest(tb, pc) return workspace.NewManager() } func registerForTest(tb testing.TB, pc poolConfig) { tb.Helper() if pc.Addr == "local" { registerLocalForTest(tb, tb.TempDir()) return } host, grpcPort := parseHostPort(pc.Addr) ports := tai.Ports{GRPC: grpcPort} res, err := tai.DialRemote(host, ports) if err != nil { tb.Fatalf("DialRemote(%s): %v", pc.Addr, err) } taiID := taiIDFromAddr(pc.Addr) reg := registry.Global() reg.Register(®istry.TaiNode{TaiID: taiID, Mode: "direct"}) reg.SetResources(taiID, res) tb.Cleanup(func() { res.Close() }) } func registerLocalForTest(tb testing.TB, dataDir string) { tb.Helper() vol := volume.NewLocal(dataDir) res, err := tai.DialLocal("", dataDir, vol) if err != nil { tb.Fatalf("DialLocal: %v", err) } reg := registry.Global() reg.Register(®istry.TaiNode{TaiID: "local", Mode: "local"}) reg.SetResources("local", res) tb.Cleanup(func() { res.Close() }) } func parseHostPort(addr string) (string, int) { addr = strings.TrimPrefix(addr, "tai://") parts := strings.SplitN(addr, ":", 2) h := parts[0] if len(parts) == 2 { if p, err := strconv.Atoi(parts[1]); err == nil { return h, p } } return h, 19100 } func setupManagerMultiNode(t *testing.T) (*workspace.Manager, string, string) { t.Helper() ensureRegistry(t) dir1 := t.TempDir() vol1 := volume.NewLocal(dir1) res1, err := tai.DialLocal("", dir1, vol1) if err != nil { t.Fatalf("DialLocal node-a: %v", err) } reg := registry.Global() reg.Register(®istry.TaiNode{TaiID: "node-a", Mode: "local"}) reg.SetResources("node-a", res1) t.Cleanup(func() { res1.Close() }) dir2 := t.TempDir() vol2 := volume.NewLocal(dir2) res2, err := tai.DialLocal("", dir2, vol2) if err != nil { t.Fatalf("DialLocal node-b: %v", err) } reg.Register(®istry.TaiNode{TaiID: "node-b", Mode: "local"}) reg.SetResources("node-b", res2) t.Cleanup(func() { res2.Close() }) return workspace.NewManager(), "node-a", "node-b" } func createWorkspace(tb testing.TB, m *workspace.Manager, node string, opts ...func(*workspace.CreateOptions)) *workspace.Workspace { tb.Helper() co := workspace.CreateOptions{ Name: "test-workspace", Owner: "test-user", Node: node, } for _, fn := range opts { fn(&co) } ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() ws, err := m.Create(ctx, co) if err != nil { tb.Fatalf("Create workspace: %v", err) } tb.Cleanup(func() { m.Delete(context.Background(), ws.ID, true) }) return ws }