ThothII/tools/tht/internal/modelmigration/migration_test.go

package modelmigration

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

	"github.com/aritmolab/thothii/tools/tht/internal/config"
	"gopkg.in/yaml.v3"
)

func TestRunEmitsValidatedCandidateWithoutChangingLegacyInputs(t *testing.T) {
	fixture := migrationFixture(t, false)
	original, _ := os.ReadFile(fixture.legacy)
	request := fixture.request()
	if err := Run(request); err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	after, _ := os.ReadFile(fixture.legacy)
	if string(after) != string(original) {
		t.Fatal("legacy installation descriptor changed")
	}
	// The strict v2 loader deliberately refuses legacy Pi sources; removing them is the
	// operator's explicit cut-over after reviewing the candidate.
	var raw legacyDescriptor
	decodeYAMLFile(t, fixture.legacy, &raw)
	if err := os.Remove(filepath.Join(raw.ProjectDirectory, "deploy", "pi", "models.json")); err != nil {
		t.Fatal(err)
	}
	if err := os.Remove(filepath.Join(raw.ProjectDirectory, "deploy", "pi", "settings.json")); err != nil {
		t.Fatal(err)
	}
	candidate, err := config.Load(request.OutputPath)
	if err != nil {
		t.Fatalf("candidate config.Load() error = %v", err)
	}
	if candidate.SchemaVersion != 2 || candidate.ModelCatalog.Defaults.Session != "zai/glm-5.3" {
		t.Fatalf("candidate = %#v", candidate.ModelCatalog)
	}
	if got := candidate.ModelCatalog.Defaults.MetadataGeneration; got != "zai/glm-5.3" {
		t.Fatalf("metadata default = %q", got)
	}
	if candidate.ModelCatalog.Providers["deepseek"].Session.Mode != "pi_builtin" {
		t.Fatal("built-in Pi model was not migrated")
	}
}

func decodeYAMLFile(t *testing.T, path string, target any) {
	t.Helper()
	source, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if err := yaml.Unmarshal(source, target); err != nil {
		t.Fatal(err)
	}
}

func TestRunReportsProviderAuthenticationConflictAndPublishesNothing(t *testing.T) {
	fixture := migrationFixture(t, true)
	err := Run(fixture.request())
	if err == nil || !strings.Contains(err.Error(), "metadataGeneration.models[0]") ||
		!strings.Contains(err.Error(), "authentication") {
		t.Fatalf("Run() error = %v", err)
	}
	if _, statErr := os.Stat(fixture.request().OutputPath); !os.IsNotExist(statErr) {
		t.Fatalf("candidate was published after conflict: %v", statErr)
	}
}

type migrationTestFixture struct {
	legacy string
	output string
}

func (f migrationTestFixture) request() Request {
	return Request{
		InstallationPath:    f.legacy,
		OutputPath:          f.output,
		SessionDefault:      "zai/glm-5.3",
		EmbeddingID:         "ollama/qwen3-embedding:0.6b",
		EmbeddingDimensions: 1024,
	}
}

func migrationFixture(t *testing.T, conflict bool) migrationTestFixture {
	t.Helper()
	base, err := filepath.EvalSymlinks(os.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	root, err := os.MkdirTemp(base, "tht-model-migration-")
	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", "pi"), 0o755); err != nil {
		t.Fatal(err)
	}
	for _, path := range []string{filepath.Join(project, "compose.yaml"), filepath.Join(project, "deploy", "compose.local.yaml")} {
		if err := os.WriteFile(path, []byte("services: {}\n"), 0o600); err != nil {
			t.Fatal(err)
		}
	}
	env := filepath.Join(root, "operator.env")
	secrets := filepath.Join(root, "secrets.env")
	auth := filepath.Join(root, "auth")
	if err := os.WriteFile(secrets, []byte("ZAI_API_KEY=test\nDEEPSEEK_API_KEY=test\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(env, []byte("THT_SECRETS_FILE="+secrets+"\nTHT_AUTH_CONFIG_ROOT="+auth+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.Mkdir(auth, 0o700); err != nil {
		t.Fatal(err)
	}
	metadata := `  default: glm-53
  models:
    - id: glm-53
      label: GLM 5.3
      litellm:
        provider: openai
        model: glm-5.3
        endpoint:
          baseUrl: https://api.z.ai/v1
      apiKeyEnv: ZAI_API_KEY
`
	if conflict {
		metadata = `  default: deepseek-v4-pro
  models:
    - id: deepseek-v4-pro
      label: DeepSeek
      litellm: {provider: deepseek, model: deepseek-v4-pro}
      apiKeyEnv: DEEPSEEK_API_KEY
`
	}
	legacy := filepath.Join(root, "legacy", "thothii-installation.yaml")
	if err := os.MkdirAll(filepath.Dir(legacy), 0o700); err != nil {
		t.Fatal(err)
	}
	source := "profile: local\nprojectDirectory: " + project + "\nenvFile: " + env +
		"\nmetadataGeneration:\n" + metadata +
		"authentication:\n  configDirectory: " + auth + "\n"
	if err := os.WriteFile(legacy, []byte(source), 0o600); err != nil {
		t.Fatal(err)
	}
	models := `{"providers":{"zai":{"baseUrl":"https://api.z.ai/v1","api":"openai-completions","apiKey":"$ZAI_API_KEY","models":[{"id":"glm-5.3","name":"GLM 5.3","reasoning":true,"contextWindow":200000,"maxTokens":131072}]}}}`
	settings := `{"defaultProjectTrust":"always","enabledModels":["deepseek/deepseek-v4-pro","zai/glm-5.3"]}`
	if err := os.WriteFile(filepath.Join(project, "deploy", "pi", "models.json"), []byte(models), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(project, "deploy", "pi", "settings.json"), []byte(settings), 0o600); err != nil {
		t.Fatal(err)
	}
	return migrationTestFixture{legacy: legacy, output: filepath.Join(root, "candidate", "thothii-installation.yaml")}
}