ThothII/tools/thothctl/internal/workspaceops/operations_test.go

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":[]}`
}

func TestParseVectorInspect(t *testing.T) {
	req, err := Parse([]string{"vector", "inspect", "--workspace", "psd", "--json"})
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	r, ok := req.(VectorInspectRequest)
	if !ok {
		t.Fatalf("got %T", req)
	}
	if r.Workspace != "psd" || !r.JSON {
		t.Fatalf("unexpected request: %+v", r)
	}
	env, err := req.stdinEnvelope()
	if err != nil {
		t.Fatalf("envelope: %v", err)
	}
	if env.WorkspaceID != "psd" || env.Collection != "" {
		t.Fatalf("unexpected envelope: %+v", env)
	}
}

func TestParseVectorRebuildGuards(t *testing.T) {
	req, err := Parse([]string{"vector", "rebuild", "--workspace", "psd", "--collection", "psd", "--confirm", "psd", "--destroy"})
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	_, ok := req.(VectorRebuildRequest)
	if !ok {
		t.Fatalf("got %T", req)
	}
	env, err := req.stdinEnvelope()
	if err != nil {
		t.Fatalf("envelope: %v", err)
	}
	if env.Collection != "psd" || !env.Destroy {
		t.Fatalf("unexpected envelope: %+v", env)
	}
}

func TestParseVectorRebuildRefusesMismatchedConfirmation(t *testing.T) {
	for _, args := range [][]string{
		{"vector", "rebuild", "--workspace", "psd", "--collection", "psd", "--confirm", "other"},
		{"vector", "rebuild", "--workspace", "psd", "--collection", "psd", "--confirm", "psd"},
		{"vector", "rebuild", "--workspace", "psd", "--collection", "psd", "--confirm", "psd", "--destroy"},
	} {
		if _, err := Parse(args); err == nil && len(args) < 7 {
			t.Fatalf("expected error for %v", args)
		}
	}
}

func TestParseVectorRebuildRequiresDestroy(t *testing.T) {
	req, err := Parse([]string{"vector", "rebuild", "--workspace", "psd", "--collection", "psd", "--confirm", "psd"})
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	if _, err := req.stdinEnvelope(); err == nil {
		t.Fatal("expected envelope error without --destroy")
	}
}

func TestParseVectorUnknownSubcommand(t *testing.T) {
	if _, err := Parse([]string{"vector", "drop", "--workspace", "psd"}); err == nil {
		t.Fatal("expected error for unknown vector command")
	}
}