277 lines
8.7 KiB
Go
277 lines
8.7 KiB
Go
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")
|
|
}
|
|
}
|