fetch_ml/tests/unit/experiment/manager_test.go

package experiment

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

	"github.com/jfraeys/fetch_ml/internal/experiment"
)

func findArchivedExperiment(t *testing.T, basePath, commitID string) string {
	t.Helper()
	archiveRoot := filepath.Join(basePath, "archive")
	var found string
	_ = filepath.WalkDir(archiveRoot, func(path string, d os.DirEntry, err error) error {
		if err != nil {
			return nil
		}
		if d.IsDir() && filepath.Base(path) == commitID {
			found = path
			return filepath.SkipDir
		}
		return nil
	})
	return found
}

const (
	experimentsPath = "/experiments"
	testCommitID    = "abc123"
)

func TestNewManager(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	// Test that manager was created successfully by checking it can generate paths
	path := manager.GetExperimentPath("test")
	if path == "" {
		t.Error("Manager should be able to generate paths")
	}
}

func TestGetExperimentPath(t *testing.T) {
	t.Parallel() // Enable parallel execution

	const experimentsPath = "/experiments"
	const testCommitID = "abc123"

	basePath := experimentsPath
	manager := experiment.NewManager(basePath)
	commitID := testCommitID

	expectedPath := filepath.Join(basePath, commitID)
	actualPath := manager.GetExperimentPath(commitID)

	if actualPath != expectedPath {
		t.Errorf("Expected path %s, got %s", expectedPath, actualPath)
	}
}

func TestGetFilesPath(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := experimentsPath
	manager := experiment.NewManager(basePath)
	commitID := testCommitID

	expectedPath := filepath.Join(basePath, commitID, "files")
	actualPath := manager.GetFilesPath(commitID)

	if actualPath != expectedPath {
		t.Errorf("Expected path %s, got %s", expectedPath, actualPath)
	}
}

func TestGetMetadataPath(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := experimentsPath
	manager := experiment.NewManager(basePath)
	commitID := testCommitID

	expectedPath := filepath.Join(basePath, commitID, "meta.bin")
	actualPath := manager.GetMetadataPath(commitID)

	if actualPath != expectedPath {
		t.Errorf("Expected path %s, got %s", expectedPath, actualPath)
	}
}

func TestExperimentExists(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	// Test non-existent experiment
	if manager.ExperimentExists("nonexistent") {
		t.Error("Experiment should not exist")
	}

	// Create experiment directory
	commitID := testCommitID
	experimentPath := manager.GetExperimentPath(commitID)
	err := os.MkdirAll(experimentPath, 0750)
	if err != nil {
		t.Fatalf("Failed to create experiment directory: %v", err)
	}

	// Test existing experiment
	if !manager.ExperimentExists(commitID) {
		t.Error("Experiment should exist")
	}
}

func TestCreateExperiment(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	commitID := testCommitID

	err := manager.CreateExperiment(commitID)
	if err != nil {
		t.Fatalf("Failed to create experiment: %v", err)
	}

	// Verify experiment directory exists
	if !manager.ExperimentExists(commitID) {
		t.Error("Experiment should exist after creation")
	}

	// Verify files directory exists
	filesPath := manager.GetFilesPath(commitID)
	info, err := os.Stat(filesPath)
	if err != nil {
		t.Fatalf("Files directory should exist: %v", err)
	}

	if !info.IsDir() {
		t.Error("Files path should be a directory")
	}
}

func TestWriteAndReadMetadata(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	commitID := testCommitID
	originalMetadata := &experiment.Metadata{
		CommitID:  commitID,
		Timestamp: time.Now().Unix(),
		JobName:   "test_experiment",
		User:      "testuser",
	}

	// Create experiment first
	err := manager.CreateExperiment(commitID)
	if err != nil {
		t.Fatalf("Failed to create experiment: %v", err)
	}

	// Write metadata
	err = manager.WriteMetadata(originalMetadata)
	if err != nil {
		t.Fatalf("Failed to write metadata: %v", err)
	}

	// Read metadata
	loadedMetadata, err := manager.ReadMetadata(commitID)
	if err != nil {
		t.Fatalf("Failed to read metadata: %v", err)
	}

	// Verify metadata
	if loadedMetadata.CommitID != originalMetadata.CommitID {
		t.Errorf("Expected commit ID %s, got %s", originalMetadata.CommitID, loadedMetadata.CommitID)
	}

	if loadedMetadata.Timestamp != originalMetadata.Timestamp {
		t.Errorf("Expected timestamp %d, got %d", originalMetadata.Timestamp, loadedMetadata.Timestamp)
	}

	if loadedMetadata.JobName != originalMetadata.JobName {
		t.Errorf("Expected job name %s, got %s", originalMetadata.JobName, loadedMetadata.JobName)
	}

	if loadedMetadata.User != originalMetadata.User {
		t.Errorf("Expected user %s, got %s", originalMetadata.User, loadedMetadata.User)
	}
}

func TestReadMetadataNonExistent(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	// Try to read metadata from non-existent experiment
	_, err := manager.ReadMetadata("nonexistent")
	if err == nil {
		t.Error("Expected error when reading metadata from non-existent experiment")
	}
}

func TestWriteMetadataNonExistentDir(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	commitID := testCommitID
	metadata := &experiment.Metadata{
		CommitID:  commitID,
		Timestamp: time.Now().Unix(),
		JobName:   "test_experiment",
		User:      "testuser",
	}

	// Try to write metadata without creating experiment directory first
	err := manager.WriteMetadata(metadata)
	if err == nil {
		t.Error("Expected error when writing metadata to non-existent experiment")
	}
}

func TestListExperiments(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	// Create multiple experiments
	experiments := []string{"abc123", "def456", "ghi789"}
	for _, commitID := range experiments {
		err := manager.CreateExperiment(commitID)
		if err != nil {
			t.Fatalf("Failed to create experiment %s: %v", commitID, err)
		}
	}

	// List experiments
	experimentList, err := manager.ListExperiments()
	if err != nil {
		t.Fatalf("Failed to list experiments: %v", err)
	}

	if len(experimentList) != 3 {
		t.Errorf("Expected 3 experiments, got %d", len(experimentList))
	}

	// Verify all experiments are listed
	for _, commitID := range experiments {
		found := false
		for _, exp := range experimentList {
			if exp == commitID {
				found = true
				break
			}
		}
		if !found {
			t.Errorf("Experiment %s not found in list", commitID)
		}
	}
}

func TestPruneExperiments(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	// Create experiments with different timestamps
	now := time.Now()
	experiments := []struct {
		commitID  string
		timestamp int64
	}{
		{"old1", now.AddDate(0, 0, -10).Unix()},
		{"old2", now.AddDate(0, 0, -5).Unix()},
		{"recent", now.AddDate(0, 0, -1).Unix()},
	}

	for _, exp := range experiments {
		// Create experiment directory
		err := manager.CreateExperiment(exp.commitID)
		if err != nil {
			t.Fatalf("Failed to create experiment %s: %v", exp.commitID, err)
		}

		// Write metadata
		metadata := &experiment.Metadata{
			CommitID:  exp.commitID,
			Timestamp: exp.timestamp,
			JobName:   "experiment_" + exp.commitID,
			User:      "testuser",
		}

		err = manager.WriteMetadata(metadata)
		if err != nil {
			t.Fatalf("Failed to write metadata for %s: %v", exp.commitID, err)
		}
	}

	// Prune experiments (keep 1, prune older than 3 days)
	pruned, err := manager.PruneExperiments(1, 3)
	if err != nil {
		t.Fatalf("Failed to prune experiments: %v", err)
	}

	// Should prune old1 and old2 (older than 3 days)
	if len(pruned) != 2 {
		t.Errorf("Expected 2 pruned experiments, got %d", len(pruned))
	}

	// Verify recent experiment still exists
	if !manager.ExperimentExists("recent") {
		t.Error("Recent experiment should still exist")
	}

	// Verify old experiments are gone
	if manager.ExperimentExists("old1") {
		t.Error("Old experiment 1 should be pruned")
	}

	if manager.ExperimentExists("old2") {
		t.Error("Old experiment 2 should be pruned")
	}

	if archived := findArchivedExperiment(t, basePath, "old1"); archived == "" {
		t.Error("Old experiment 1 should be archived")
	}

	if archived := findArchivedExperiment(t, basePath, "old2"); archived == "" {
		t.Error("Old experiment 2 should be archived")
	}
}

func TestPruneExperimentsKeepCount(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	// Create experiments with different timestamps
	now := time.Now()
	experiments := []string{"exp1", "exp2", "exp3", "exp4"}

	for i, commitID := range experiments {
		// Create experiment directory
		err := manager.CreateExperiment(commitID)
		if err != nil {
			t.Fatalf("Failed to create experiment %s: %v", commitID, err)
		}

		// Write metadata with different timestamps (newer first)
		metadata := &experiment.Metadata{
			CommitID:  commitID,
			Timestamp: now.Add(-time.Duration(i) * time.Hour).Unix(),
			JobName:   "experiment_" + commitID,
			User:      "testuser",
		}

		err = manager.WriteMetadata(metadata)
		if err != nil {
			t.Fatalf("Failed to write metadata for %s: %v", commitID, err)
		}
	}

	// Prune experiments (keep 2 newest, no age limit)
	pruned, err := manager.PruneExperiments(2, 0)
	if err != nil {
		t.Fatalf("Failed to prune experiments: %v", err)
	}

	// Should prune 2 oldest experiments
	if len(pruned) != 2 {
		t.Errorf("Expected 2 pruned experiments, got %d", len(pruned))
	}

	// Verify newest experiments still exist
	if !manager.ExperimentExists("exp1") {
		t.Error("Newest experiment should still exist")
	}

	if !manager.ExperimentExists("exp2") {
		t.Error("Second newest experiment should still exist")
	}

	// Verify oldest experiments are gone
	if manager.ExperimentExists("exp3") {
		t.Error("Old experiment 3 should be pruned")
	}

	if manager.ExperimentExists("exp4") {
		t.Error("Old experiment 4 should be pruned")
	}

	if archived := findArchivedExperiment(t, basePath, "exp3"); archived == "" {
		t.Error("Old experiment 3 should be archived")
	}

	if archived := findArchivedExperiment(t, basePath, "exp4"); archived == "" {
		t.Error("Old experiment 4 should be archived")
	}
}

func TestMetadataPartialFields(t *testing.T) {
	t.Parallel() // Enable parallel execution

	basePath := t.TempDir()
	manager := experiment.NewManager(basePath)

	commitID := "abc123"

	// Create experiment
	err := manager.CreateExperiment(commitID)
	if err != nil {
		t.Fatalf("Failed to create experiment: %v", err)
	}

	// Test metadata with only required fields
	metadata := &experiment.Metadata{
		CommitID:  commitID,
		Timestamp: time.Now().Unix(),
		// JobName and User are optional
	}

	err = manager.WriteMetadata(metadata)
	if err != nil {
		t.Fatalf("Failed to write metadata: %v", err)
	}

	// Read it back
	loadedMetadata, err := manager.ReadMetadata(commitID)
	if err != nil {
		t.Fatalf("Failed to read metadata: %v", err)
	}

	if loadedMetadata.CommitID != commitID {
		t.Errorf("Expected commit ID %s, got %s", commitID, loadedMetadata.CommitID)
	}

	if loadedMetadata.JobName != "" {
		t.Errorf("Expected empty job name, got %s", loadedMetadata.JobName)
	}

	if loadedMetadata.User != "" {
		t.Errorf("Expected empty user, got %s", loadedMetadata.User)
	}
}