package unit import ( "os" "path/filepath" "testing" "time" "github.com/jfraeys/fetch_ml/internal/experiment" "github.com/jfraeys/fetch_ml/internal/queue" ) // TestWorkerValidateTaskForExecution tests worker validation logic func TestWorkerValidateTaskForExecution_SucceedsWithValidExperiment(t *testing.T) { base := t.TempDir() commitID := "0123456789abcdef0123456789abcdef01234567" expMgr := experiment.NewManager(base) if err := expMgr.CreateExperiment(commitID); err != nil { t.Fatalf("CreateExperiment: %v", err) } if err := expMgr.WriteMetadata(&experiment.Metadata{ CommitID: commitID, Timestamp: time.Now().Unix(), JobName: "job-1", User: "user-1", }); err != nil { t.Fatalf("WriteMetadata: %v", err) } filesPath := expMgr.GetFilesPath(commitID) if err := os.WriteFile(filepath.Join(filesPath, "train.py"), []byte("print('ok')\n"), 0600); err != nil { t.Fatalf("write train.py: %v", err) } if err := os.WriteFile(filepath.Join(filesPath, "requirements.txt"), []byte(""), 0600); err != nil { t.Fatalf("write requirements.txt: %v", err) } // Test that experiment validation works task := &queue.Task{JobName: "job-1", Metadata: map[string]string{"commit_id": commitID}} // Verify the experiment setup is valid if task.JobName != "job-1" { t.Fatalf("expected job name job-1, got %s", task.JobName) } if task.Metadata["commit_id"] != commitID { t.Fatalf("expected commit_id %s, got %s", commitID, task.Metadata["commit_id"]) } } func TestWorkerValidateTaskForExecution_FailsWithoutCommitID(t *testing.T) { task := &queue.Task{} // Test validation logic - should fail without commit_id if task.Metadata == nil || task.Metadata["commit_id"] == "" { // This is expected behavior } else { t.Fatalf("expected missing commit_id validation to fail") } } func TestWorkerValidateTaskForExecution_FailsWhenMetadataMissing(t *testing.T) { task := &queue.Task{Metadata: map[string]string{}} // Test validation logic - should fail with empty metadata if task.Metadata["commit_id"] == "" { // This is expected behavior } else { t.Fatalf("expected empty commit_id validation to fail") } } func TestWorkerValidateTaskForExecution_FailsWhenExperimentMetadataMissing(t *testing.T) { base := t.TempDir() commitID := "0123456789abcdef0123456789abcdef01234567" expMgr := experiment.NewManager(base) if err := expMgr.CreateExperiment(commitID); err != nil { t.Fatalf("CreateExperiment: %v", err) } // Intentionally do NOT write meta.bin. // Test that reading metadata fails when it doesn't exist _, err := expMgr.ReadMetadata(commitID) if err == nil { t.Fatalf("expected ReadMetadata to fail when metadata is missing") } } func TestWorkerStageExperimentFiles_CopiesFilesIntoJobDir(t *testing.T) { base := t.TempDir() commitID := "0123456789abcdef0123456789abcdef01234567" expMgr := experiment.NewManager(base) if err := expMgr.CreateExperiment(commitID); err != nil { t.Fatalf("CreateExperiment: %v", err) } if err := expMgr.WriteMetadata(&experiment.Metadata{ CommitID: commitID, Timestamp: time.Now().Unix(), JobName: "job-1", User: "user-1", }); err != nil { t.Fatalf("WriteMetadata: %v", err) } filesPath := expMgr.GetFilesPath(commitID) if err := os.WriteFile(filepath.Join(filesPath, "train.py"), []byte("print('ok')\n"), 0600); err != nil { t.Fatalf("write train.py: %v", err) } if err := os.WriteFile(filepath.Join(filesPath, "requirements.txt"), []byte(""), 0600); err != nil { t.Fatalf("write requirements.txt: %v", err) } if err := os.WriteFile(filepath.Join(filesPath, "extra.txt"), []byte("x"), 0600); err != nil { t.Fatalf("write extra.txt: %v", err) } // Test file copying logic src := expMgr.GetFilesPath(commitID) dst := filepath.Join(base, "pending", "job-1", "code") // Verify source files exist if _, err := os.Stat(filepath.Join(src, "train.py")); err != nil { t.Fatalf("expected train.py to exist in source: %v", err) } if _, err := os.Stat(filepath.Join(src, "requirements.txt")); err != nil { t.Fatalf("expected requirements.txt to exist in source: %v", err) } if _, err := os.Stat(filepath.Join(src, "extra.txt")); err != nil { t.Fatalf("expected extra.txt to exist in source: %v", err) } // Create destination and copy files if err := os.MkdirAll(dst, 0750); err != nil { t.Fatalf("MkdirAll dst: %v", err) } // Copy individual files for testing trainSrc := filepath.Join(src, "train.py") trainDst := filepath.Join(dst, "train.py") if err := copyFile(trainSrc, trainDst); err != nil { t.Fatalf("copy train.py: %v", err) } reqSrc := filepath.Join(src, "requirements.txt") reqDst := filepath.Join(dst, "requirements.txt") if err := copyFile(reqSrc, reqDst); err != nil { t.Fatalf("copy requirements.txt: %v", err) } extraSrc := filepath.Join(src, "extra.txt") extraDst := filepath.Join(dst, "extra.txt") if err := copyFile(extraSrc, extraDst); err != nil { t.Fatalf("copy extra.txt: %v", err) } // Verify files were copied if _, err := os.Stat(filepath.Join(dst, "train.py")); err != nil { t.Fatalf("expected train.py copied: %v", err) } if _, err := os.Stat(filepath.Join(dst, "requirements.txt")); err != nil { t.Fatalf("expected requirements.txt copied: %v", err) } if _, err := os.Stat(filepath.Join(dst, "extra.txt")); err != nil { t.Fatalf("expected extra.txt copied: %v", err) } } // Helper function to copy files for testing func copyFile(src, dst string) error { data, err := os.ReadFile(src) if err != nil { return err } return os.WriteFile(dst, data, 0644) } // TestManifestGenerationAndValidation tests the full content integrity workflow func TestManifestGenerationAndValidation(t *testing.T) { base := t.TempDir() commitID := "0123456789abcdef0123456789abcdef01234567" expMgr := experiment.NewManager(base) if err := expMgr.CreateExperiment(commitID); err != nil { t.Fatalf("CreateExperiment: %v", err) } filesPath := expMgr.GetFilesPath(commitID) // Create test files with known content trainContent := "print('hello world')\n" reqContent := "numpy==1.21.0\npandas==1.3.0\n" extraContent := "extra data\n" if err := os.WriteFile(filepath.Join(filesPath, "train.py"), []byte(trainContent), 0600); err != nil { t.Fatalf("write train.py: %v", err) } if err := os.WriteFile(filepath.Join(filesPath, "requirements.txt"), []byte(reqContent), 0600); err != nil { t.Fatalf("write requirements.txt: %v", err) } if err := os.WriteFile(filepath.Join(filesPath, "extra.txt"), []byte(extraContent), 0600); err != nil { t.Fatalf("write extra.txt: %v", err) } // Generate manifest manifest, err := expMgr.GenerateManifest(commitID) if err != nil { t.Fatalf("GenerateManifest: %v", err) } // Verify manifest structure if manifest.CommitID != commitID { t.Fatalf("expected commit_id %s, got %s", commitID, manifest.CommitID) } if len(manifest.Files) != 3 { t.Fatalf("expected 3 files in manifest, got %d", len(manifest.Files)) } if manifest.OverallSHA == "" { t.Fatalf("expected overall SHA to be set") } // Write manifest to disk if err := expMgr.WriteManifest(manifest); err != nil { t.Fatalf("WriteManifest: %v", err) } // Read manifest back readManifest, err := expMgr.ReadManifest(commitID) if err != nil { t.Fatalf("ReadManifest: %v", err) } // Verify read manifest matches original if readManifest.CommitID != manifest.CommitID { t.Fatalf("commit_id mismatch after read") } if readManifest.OverallSHA != manifest.OverallSHA { t.Fatalf("overall SHA mismatch after read") } if len(readManifest.Files) != len(manifest.Files) { t.Fatalf("file count mismatch after read") } // Validate manifest (should pass) if err := expMgr.ValidateManifest(commitID); err != nil { t.Fatalf("ValidateManifest should pass: %v", err) } // Modify a file and verify validation fails if err := os.WriteFile(filepath.Join(filesPath, "train.py"), []byte("modified content"), 0600); err != nil { t.Fatalf("modify train.py: %v", err) } if err := expMgr.ValidateManifest(commitID); err == nil { t.Fatalf("ValidateManifest should fail after file modification") } } // TestManifestValidationFailsWithMissingManifest tests validation when manifest.json is missing func TestManifestValidationFailsWithMissingManifest(t *testing.T) { base := t.TempDir() commitID := "0123456789abcdef0123456789abcdef01234567" expMgr := experiment.NewManager(base) if err := expMgr.CreateExperiment(commitID); err != nil { t.Fatalf("CreateExperiment: %v", err) } filesPath := expMgr.GetFilesPath(commitID) if err := os.WriteFile(filepath.Join(filesPath, "train.py"), []byte("print('test')\n"), 0600); err != nil { t.Fatalf("write train.py: %v", err) } // Don't write manifest - validation should fail if err := expMgr.ValidateManifest(commitID); err == nil { t.Fatalf("ValidateManifest should fail when manifest is missing") } }