package workspaceops import ( "bytes" "context" "encoding/json" "fmt" "io" "os" "path/filepath" "reflect" "strings" "testing" "github.com/aritmolab/thothii/tools/thothctl/internal/compose" "github.com/aritmolab/thothii/tools/thothctl/internal/config" ) type fakeRunner struct { run func(args []string, stdin string) (compose.Result, error) all [][]string stdins []string } func (r *fakeRunner) Run(_ context.Context, args []string, stdin io.Reader) (compose.Result, error) { payload := "" if stdin != nil { bytes, err := io.ReadAll(stdin) if err != nil { return compose.Result{}, err } payload = string(bytes) } r.all = append(r.all, append([]string(nil), args...)) r.stdins = append(r.stdins, payload) return r.run(args, payload) } func TestExecuteSuggestFksStreamsSQLFileContentsOnStdin(t *testing.T) { installation := testInstallation(t) sqlPath := filepath.Join(filepath.Dir(installation.Path), "query.sql") if err := os.WriteFile(sqlPath, []byte("select 1;\n"), 0o600); err != nil { t.Fatal(err) } runner := &fakeRunner{run: func(args []string, stdin string) (compose.Result, error) { switch { case contains(args, "config", "--format", "json"): return compose.Result{Stdout: `{"services":{"core":{"image":"thothii-core:local"}}}`}, nil case reflect.DeepEqual(args, []string{"image", "inspect", "--format", "{{.Id}}", "thothii-core:local"}): return compose.Result{Stdout: "sha256:" + strings.Repeat("a", 64)}, nil case contains(args, "workspace-maintenance", "schema-suggest-fks"): var envelope map[string]any if err := json.Unmarshal([]byte(stdin), &envelope); err != nil { t.Fatalf("stdin JSON = %q, err=%v", stdin, err) } sqlFiles, ok := envelope["fromSql"].([]any) if !ok || len(sqlFiles) != 1 { t.Fatalf("sqlFiles = %#v", envelope["fromSql"]) } file, ok := sqlFiles[0].(map[string]any) if !ok || file["sql"] != "select 1;\n" || file["name"] != "query.sql" { t.Fatalf("sql file envelope = %#v", sqlFiles[0]) } return compose.Result{Stdout: successResult("schema-suggest-fks")}, nil default: t.Fatalf("unexpected Docker invocation: %#v", args) return compose.Result{}, nil } }} _, err := Execute(context.Background(), installation, runner, SuggestFksRequest{baseRequest: baseRequest{Workspace: "abc"}, FromSQL: []string{sqlPath}}) if err != nil { t.Fatalf("Execute() error = %v", err) } } func TestExecuteSuggestFksRejectsTotalSQLIngressOverSixteenMiB(t *testing.T) { installation := testInstallation(t) paths := make([]string, 0, 17) for index := 0; index < 17; index++ { path := filepath.Join(filepath.Dir(installation.Path), fmt.Sprintf("query-%02d.sql", index)) if err := os.WriteFile(path, bytes.Repeat([]byte("x"), 1<<20), 0o600); err != nil { t.Fatal(err) } paths = append(paths, path) } runner := &fakeRunner{run: func(args []string, stdin string) (compose.Result, error) { t.Fatalf("Docker should not run when total SQL ingress exceeds the bound: %#v", args) return compose.Result{}, nil }} _, err := Execute(context.Background(), installation, runner, SuggestFksRequest{baseRequest: baseRequest{Workspace: "abc"}, FromSQL: paths}) if err == nil || !strings.Contains(err.Error(), "total") { t.Fatalf("Execute() error = %v, want total-size failure", err) } } func TestExecuteSchemaCheckStreamsAnnotationContentOnStdin(t *testing.T) { installation := testInstallation(t) annotationsPath := filepath.Join(filepath.Dir(installation.Path), "annotations.yaml") if err := os.WriteFile(annotationsPath, []byte("reviewed: []\n"), 0o600); err != nil { t.Fatal(err) } runner := &fakeRunner{run: func(args []string, stdin string) (compose.Result, error) { switch { case contains(args, "config", "--format", "json"): return compose.Result{Stdout: `{"services":{"core":{"image":"thothii-core:local"}}}`}, nil case reflect.DeepEqual(args, []string{"image", "inspect", "--format", "{{.Id}}", "thothii-core:local"}): return compose.Result{Stdout: "sha256:" + strings.Repeat("b", 64)}, nil case contains(args, "workspace-maintenance", "schema-check"): var envelope map[string]any if err := json.Unmarshal([]byte(stdin), &envelope); err != nil { t.Fatalf("stdin JSON = %q, err=%v", stdin, err) } if envelope["annotationsYaml"] != "reviewed: []\n" || envelope["reviewedCandidatesDigest"] != "sha256:"+strings.Repeat("c", 64) { t.Fatalf("annotation envelope = %#v", envelope) } if _, exists := envelope["annotationsPath"]; exists { t.Fatalf("annotation path leaked into stdin: %#v", envelope) } return compose.Result{Stdout: successResult("schema-check")}, nil default: t.Fatalf("unexpected Docker invocation: %#v", args) return compose.Result{}, nil } }} _, err := Execute(context.Background(), installation, runner, CheckSchemaRequest{baseRequest: baseRequest{Workspace: "abc"}, Annotations: annotationsPath, ReviewedCandidates: "sha256:" + strings.Repeat("c", 64)}) if err != nil { t.Fatalf("Execute() error = %v", err) } } func TestExecuteSuggestFksWritesTheReturnedCandidateArtifact(t *testing.T) { installation := testInstallation(t) outputPath := filepath.Join(filepath.Dir(installation.Path), "candidates.yaml") runner := &fakeRunner{run: func(args []string, stdin string) (compose.Result, error) { switch { case contains(args, "config", "--format", "json"): return compose.Result{Stdout: `{"services":{"core":{"image":"thothii-core:local"}}}`}, nil case reflect.DeepEqual(args, []string{"image", "inspect", "--format", "{{.Id}}", "thothii-core:local"}): return compose.Result{Stdout: "sha256:" + strings.Repeat("d", 64)}, nil case contains(args, "workspace-maintenance", "schema-suggest-fks"): return compose.Result{Stdout: `{"schemaVersion":1,"status":"blocked","code":"manual_review_required","workspaceId":"abc","workspaceRevision":"1234567890abcdef1234567890abcdef12345678","descriptorBlob":"sha256:` + strings.Repeat("e", 64) + `","operation":"schema-suggest-fks","completedStages":[],"suggestedFksYaml":"reviewed: []\n"}`}, nil default: t.Fatalf("unexpected Docker invocation: %#v", args) return compose.Result{}, nil } }} _, err := Execute(context.Background(), installation, runner, SuggestFksRequest{baseRequest: baseRequest{Workspace: "abc"}, Output: outputPath}) if err != nil { t.Fatalf("Execute() error = %v", err) } contents, err := os.ReadFile(outputPath) if err != nil { t.Fatalf("output artifact missing: %v", err) } if string(contents) != "reviewed: []\n" { t.Fatalf("output contents = %q", contents) } } func testInstallation(t *testing.T) config.Installation { t.Helper() temporaryRoot, err := filepath.EvalSymlinks(os.TempDir()) if err != nil { t.Fatal(err) } root, err := os.MkdirTemp(temporaryRoot, "thothctl-workspaceops-") if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = os.RemoveAll(root) }) project := filepath.Join(root, "project") if err := os.MkdirAll(filepath.Join(project, "deploy"), 0o700); err != nil { t.Fatal(err) } if err := os.WriteFile(filepath.Join(project, "compose.yaml"), []byte("services: {}\n"), 0o600); err != nil { t.Fatal(err) } if err := os.WriteFile(filepath.Join(project, "deploy", "compose.local.yaml"), []byte("services: {}\n"), 0o600); err != nil { t.Fatal(err) } envFile := filepath.Join(root, "installation.env") if err := os.WriteFile(envFile, []byte("SAFE_VALUE=1\n"), 0o600); err != nil { t.Fatal(err) } return config.Installation{Path: filepath.Join(root, "thothii-installation.yaml"), Profile: "local", ProjectDirectory: project, EnvFile: envFile} } func contains(values []string, sequence ...string) bool { for start := range values { if start+len(sequence) <= len(values) && reflect.DeepEqual(values[start:start+len(sequence)], sequence) { return true } } return false } func successResult(operation string) string { return `{"schemaVersion":1,"status":"succeeded","code":"ok","workspaceId":"abc","workspaceRevision":"1234567890abcdef1234567890abcdef12345678","descriptorBlob":"sha256:` + strings.Repeat("f", 64) + `","operation":"` + operation + `","completedStages":[]}` }