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