fetch_ml/internal/storage/db_test.go

package storage

import (
	"os"
	"testing"
)

func TestDB(t *testing.T) {
	// Use a temporary database
	dbPath := t.TempDir() + "/test.db"

	// Initialize database
	db, err := NewDBFromPath(dbPath)
	if err != nil {
		t.Fatalf("Failed to create database: %v", err)
	}
	defer db.Close()

	// Initialize schema
	schema, err := os.ReadFile("schema.sql")
	if err != nil {
		t.Fatalf("Failed to read schema: %v", err)
	}

	if err := db.Initialize(string(schema)); err != nil {
		t.Fatalf("Failed to initialize schema: %v", err)
	}

	// Test job creation
	job := &Job{
		ID:       "test-job-1",
		JobName:  "test_experiment",
		Args:     "--epochs 10 --lr 0.001",
		Status:   "pending",
		Priority: 1,
		Datasets: []string{"dataset1", "dataset2"},
		Metadata: map[string]string{"gpu": "true", "memory": "8GB"},
	}

	if err := db.CreateJob(job); err != nil {
		t.Fatalf("Failed to create job: %v", err)
	}

	// Verify job exists in database
	var count int
	err = db.conn.QueryRow("SELECT COUNT(*) FROM jobs WHERE id = ?", "test-job-1").Scan(&count)
	if err != nil {
		t.Fatalf("Failed to verify job creation: %v", err)
	}
	if count != 1 {
		t.Fatalf("Expected 1 job in database, got %d", count)
	}

	// Test job retrieval
	retrievedJob, err := db.GetJob("test-job-1")
	if err != nil {
		t.Fatalf("Failed to get job: %v", err)
	}

	if retrievedJob.ID != job.ID {
		t.Errorf("Expected job ID %s, got %s", job.ID, retrievedJob.ID)
	}

	if retrievedJob.JobName != job.JobName {
		t.Errorf("Expected job name %s, got %s", job.JobName, retrievedJob.JobName)
	}

	if len(retrievedJob.Datasets) != 2 {
		t.Errorf("Expected 2 datasets, got %d", len(retrievedJob.Datasets))
	}

	if retrievedJob.Metadata["gpu"] != "true" {
		t.Errorf("Expected gpu=true, got %s", retrievedJob.Metadata["gpu"])
	}

	// Test job status update
	if err := db.UpdateJobStatus("test-job-1", "running", "worker-1", ""); err != nil {
		t.Fatalf("Failed to update job status: %v", err)
	}

	// Verify status update
	updatedJob, err := db.GetJob("test-job-1")
	if err != nil {
		t.Fatalf("Failed to get updated job: %v", err)
	}

	if updatedJob.Status != "running" {
		t.Errorf("Expected status running, got %s", updatedJob.Status)
	}

	if updatedJob.WorkerID != "worker-1" {
		t.Errorf("Expected worker ID worker-1, got %s", updatedJob.WorkerID)
	}

	if updatedJob.StartedAt == nil {
		t.Error("Expected StartedAt to be set")
	}

	// Test worker registration
	worker := &Worker{
		ID:          "worker-1",
		Hostname:    "test-host",
		Status:      "active",
		CurrentJobs: 0,
		MaxJobs:     2,
		Metadata:    map[string]string{"cpu": "8", "memory": "16GB"},
	}

	if err := db.RegisterWorker(worker); err != nil {
		t.Fatalf("Failed to register worker: %v", err)
	}

	// Test worker heartbeat
	if err := db.UpdateWorkerHeartbeat("worker-1"); err != nil {
		t.Fatalf("Failed to update worker heartbeat: %v", err)
	}

	// Test metrics recording
	if err := db.RecordJobMetric("test-job-1", "accuracy", "0.95"); err != nil {
		t.Fatalf("Failed to record job metric: %v", err)
	}

	if err := db.RecordSystemMetric("cpu_usage", "75"); err != nil {
		t.Fatalf("Failed to record system metric: %v", err)
	}

	// Test metrics retrieval
	metrics, err := db.GetJobMetrics("test-job-1")
	if err != nil {
		t.Fatalf("Failed to get job metrics: %v", err)
	}

	if metrics["accuracy"] != "0.95" {
		t.Errorf("Expected accuracy 0.95, got %s", metrics["accuracy"])
	}

	// Test job listing
	jobs, err := db.ListJobs("", 10)
	if err != nil {
		t.Fatalf("Failed to list jobs: %v", err)
	}

	t.Logf("Found %d jobs", len(jobs))
	for i, job := range jobs {
		t.Logf("Job %d: ID=%s, Status=%s", i, job.ID, job.Status)
	}

	if len(jobs) != 1 {
		t.Errorf("Expected 1 job, got %d", len(jobs))
		return
	}

	if jobs[0].ID != "test-job-1" {
		t.Errorf("Expected job ID test-job-1, got %s", jobs[0].ID)
		return
	}

	// Test active workers
	workers, err := db.GetActiveWorkers()
	if err != nil {
		t.Fatalf("Failed to get active workers: %v", err)
	}

	if len(workers) != 1 {
		t.Errorf("Expected 1 active worker, got %d", len(workers))
	}

	if workers[0].ID != "worker-1" {
		t.Errorf("Expected worker ID worker-1, got %s", workers[0].ID)
	}
}

func TestDBConstraints(t *testing.T) {
	dbPath := t.TempDir() + "/test_constraints.db"

	db, err := NewDBFromPath(dbPath)
	if err != nil {
		t.Fatalf("Failed to create database: %v", err)
	}
	defer db.Close()

	schema, err := os.ReadFile("schema.sql")
	if err != nil {
		t.Fatalf("Failed to read schema: %v", err)
	}

	if err := db.Initialize(string(schema)); err != nil {
		t.Fatalf("Failed to initialize schema: %v", err)
	}

	// Test duplicate job ID
	job := &Job{
		ID:      "duplicate-test",
		JobName: "test",
		Status:  "pending",
	}

	if err := db.CreateJob(job); err != nil {
		t.Fatalf("Failed to create first job: %v", err)
	}

	// Should fail on duplicate
	if err := db.CreateJob(job); err == nil {
		t.Error("Expected error when creating duplicate job")
	}

	// Test getting non-existent job
	_, err = db.GetJob("non-existent")
	if err == nil {
		t.Error("Expected error when getting non-existent job")
	}
}