fetch_ml/tests/unit/experiment/manager_test.go

package experiment

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

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

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

	basePath := "/experiments"
	manager := experiment.NewManager(basePath)
	commitID := "abc123"

	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 := "/experiments"
	manager := experiment.NewManager(basePath)
	commitID := "abc123"

	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 := "/experiments"
	manager := experiment.NewManager(basePath)
	commitID := "abc123"

	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 := "abc123"
	experimentPath := manager.GetExperimentPath(commitID)
	err := os.MkdirAll(experimentPath, 0755)
	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 := "abc123"

	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 := "abc123"
	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 := "abc123"
	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")
	}
}

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")
	}
}

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)
	}
}