fetch_ml/tests/integration/jupyter_experiment_test.go

package tests

import (
	"context"
	"log/slog"
	"os"
	"path/filepath"
	"testing"
	"time"

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

func TestJupyterExperimentIntegration(t *testing.T) {
	// Setup test environment
	tempDir, err := os.MkdirTemp("", "TestJupyterExperimentIntegration")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer func() {
		if err := os.RemoveAll(tempDir); err != nil {
			t.Logf("Warning: failed to cleanup temp dir: %v", err)
		}
	}()

	// Initialize experiment manager
	expManager := experiment.NewManager(filepath.Join(tempDir, "experiments"))
	if err := expManager.Initialize(); err != nil {
		t.Fatalf("Failed to initialize experiment manager: %v", err)
	}

	// Initialize Jupyter service manager with clean state
	serviceConfig := &jupyter.ServiceConfig{
		DefaultImage:     "test-image",
		DefaultPort:      8888,
		DefaultWorkspace: tempDir,
		MaxServices:      5,
		DefaultResources: jupyter.ResourceConfig{
			MemoryLimit: "1G",
			CPULimit:    "1",
			GPUDevices:  nil,
		},
	}

	logger := logging.NewLogger(slog.LevelInfo, false)
	serviceManager, err := jupyter.NewServiceManager(logger, serviceConfig)
	if err != nil {
		t.Fatalf("Failed to create service manager: %v", err)
	}
	defer func() {
		if err := serviceManager.Close(context.Background()); err != nil {
			t.Logf("Warning: failed to close service manager: %v", err)
		}
	}()

	// Clear any existing metadata to ensure test isolation
	if err := serviceManager.ClearAllMetadata(); err != nil {
		t.Fatalf("Failed to clear metadata: %v", err)
	}

	// Test 1: Create workspace
	workspacePath := filepath.Join(tempDir, "test_workspace")
	if err := os.MkdirAll(workspacePath, 0750); err != nil {
		t.Fatalf("Failed to create workspace: %v", err)
	}

	// Create sample files in workspace
	notebookContent := `{
		"cells": [{"cell_type": "code", "source": ["print('Hello Jupyter!')"]}],
		"metadata": {"kernelspec": {"name": "python3"}},
		"nbformat": 4
	}`
	if err := os.WriteFile(filepath.Join(workspacePath, "test.ipynb"), []byte(notebookContent), 0600); err != nil {
		t.Fatalf("Failed to create notebook: %v", err)
	}

	// Test 2: Create experiment
	experimentID := "test_experiment_123"
	if err := expManager.CreateExperiment(experimentID); err != nil {
		t.Fatalf("Failed to create experiment: %v", err)
	}

	// Write experiment metadata
	metadata := &experiment.Metadata{
		CommitID:  experimentID,
		Timestamp: time.Now().Unix(),
		JobName:   "test_job",
		User:      "test_user",
	}
	if err := expManager.WriteMetadata(metadata); err != nil {
		t.Fatalf("Failed to write experiment metadata: %v", err)
	}

	// Test 3: Link workspace with experiment
	serviceID := "test_service_456"
	if err := serviceManager.LinkWorkspaceWithExperiment(workspacePath, experimentID, serviceID); err != nil {
		t.Fatalf("Failed to link workspace with experiment: %v", err)
	}

	// Test 4: Verify workspace metadata
	workspaceMeta, err := serviceManager.GetWorkspaceMetadata(workspacePath)
	if err != nil {
		t.Fatalf("Failed to get workspace metadata: %v", err)
	}

	if workspaceMeta.ExperimentID != experimentID {
		t.Errorf("Expected experiment ID %s, got %s", experimentID, workspaceMeta.ExperimentID)
	}

	if workspaceMeta.ServiceID != serviceID {
		t.Errorf("Expected service ID %s, got %s", serviceID, workspaceMeta.ServiceID)
	}

	// Test 5: Sync workspace with experiment
	ctx := context.Background()
	if err := serviceManager.SyncWorkspaceWithExperiment(ctx, workspacePath, experimentID, "push"); err != nil {
		t.Fatalf("Failed to sync workspace: %v", err)
	}

	// Verify sync timestamp updated
	syncedMeta, err := serviceManager.GetWorkspaceMetadata(workspacePath)
	if err != nil {
		t.Fatalf("Failed to get workspace metadata after sync: %v", err)
	}

	if syncedMeta.LastSync.IsZero() {
		t.Error("Expected last sync time to be set after sync")
	}

	// Test 6: List linked workspaces
	linkedWorkspaces := serviceManager.ListLinkedWorkspaces()
	if len(linkedWorkspaces) != 1 {
		t.Errorf("Expected 1 linked workspace, got %d", len(linkedWorkspaces))
	}

	if linkedWorkspaces[0].ExperimentID != experimentID {
		t.Errorf("Expected experiment ID %s, got %s", experimentID, linkedWorkspaces[0].ExperimentID)
	}

	// Test 7: Get workspaces for experiment
	workspacesForExp := serviceManager.GetWorkspacesForExperiment(experimentID)
	if len(workspacesForExp) != 1 {
		t.Errorf("Expected 1 workspace for experiment, got %d", len(workspacesForExp))
	}

	if workspacesForExp[0].WorkspacePath != workspacePath {
		t.Errorf("Expected workspace path %s, got %s", workspacePath, workspacesForExp[0].WorkspacePath)
	}

	// Test 8: Set auto-sync
	if err := serviceManager.SetAutoSync(workspacePath, true, 15*time.Minute); err != nil {
		t.Fatalf("Failed to set auto-sync: %v", err)
	}

	// Verify auto-sync settings
	autoSyncMeta, err := serviceManager.GetWorkspaceMetadata(workspacePath)
	if err != nil {
		t.Fatalf("Failed to get workspace metadata after auto-sync: %v", err)
	}

	if !autoSyncMeta.AutoSync {
		t.Error("Expected auto-sync to be enabled")
	}

	if autoSyncMeta.SyncInterval != 15*time.Minute {
		t.Errorf("Expected sync interval 15m, got %v", autoSyncMeta.SyncInterval)
	}

	// Test 9: Add tags
	if err := serviceManager.AddTag(workspacePath, "test"); err != nil {
		t.Fatalf("Failed to add tag: %v", err)
	}

	// Verify tag added
	taggedMeta, err := serviceManager.GetWorkspaceMetadata(workspacePath)
	if err != nil {
		t.Fatalf("Failed to get workspace metadata after adding tag: %v", err)
	}

	foundTag := false
	for _, tag := range taggedMeta.Tags {
		if tag == "test" {
			foundTag = true
			break
		}
	}

	if !foundTag {
		t.Error("Expected 'test' tag to be found")
	}

	// Test 10: Unlink workspace
	if err := serviceManager.UnlinkWorkspace(workspacePath); err != nil {
		t.Fatalf("Failed to unlink workspace: %v", err)
	}

	// Verify workspace is unlinked
	_, err = serviceManager.GetWorkspaceMetadata(workspacePath)
	if err == nil {
		t.Error("Expected workspace to be unlinked")
	}

	// Test 11: Verify no linked workspaces
	linkedWorkspaces = serviceManager.ListLinkedWorkspaces()
	if len(linkedWorkspaces) != 0 {
		t.Errorf("Expected 0 linked workspaces after unlink, got %d", len(linkedWorkspaces))
	}
}