flamenco/tests/database/migration_test.go

package database_test

// SPDX-License-Identifier: GPL-3.0-or-later

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/pressly/goose/v3"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	_ "modernc.org/sqlite"

	"projects.blender.org/studio/flamenco/internal/manager/persistence"
	"projects.blender.org/studio/flamenco/pkg/api"
	"projects.blender.org/studio/flamenco/tests/helpers"
)

// DatabaseTestSuite provides comprehensive database testing
type DatabaseTestSuite struct {
	suite.Suite
	testHelper      *helpers.TestHelper
	testDBPath      string
	db              *sql.DB
	persistenceDB   *persistence.DB
}

// MigrationTestResult tracks migration test results
type MigrationTestResult struct {
	Version     int64
	Success     bool
	Duration    time.Duration
	Error       error
	Description string
}

// DataIntegrityTest represents a data integrity test case
type DataIntegrityTest struct {
	Name        string
	SetupFunc   func(*sql.DB) error
	TestFunc    func(*sql.DB) error
	CleanupFunc func(*sql.DB) error
}

// SetupSuite initializes the database test environment
func (suite *DatabaseTestSuite) SetupSuite() {
	suite.testHelper = helpers.NewTestHelper(suite.T())
	
	// Create test database
	testDir := suite.testHelper.CreateTempDir("db-tests")
	suite.testDBPath = filepath.Join(testDir, "test_flamenco.sqlite")
}

// TearDownSuite cleans up the database test environment
func (suite *DatabaseTestSuite) TearDownSuite() {
	if suite.db != nil {
		suite.db.Close()
	}
	if suite.testHelper != nil {
		suite.testHelper.Cleanup()
	}
}

// SetupTest prepares a fresh database for each test
func (suite *DatabaseTestSuite) SetupTest() {
	// Remove existing test database
	os.Remove(suite.testDBPath)
	
	// Create fresh database connection
	var err error
	suite.db, err = sql.Open("sqlite", suite.testDBPath)
	require.NoError(suite.T(), err)
	
	// Set SQLite pragmas for testing
	pragmas := []string{
		"PRAGMA foreign_keys = ON",
		"PRAGMA journal_mode = WAL",
		"PRAGMA synchronous = NORMAL",
		"PRAGMA cache_size = -64000", // 64MB cache
		"PRAGMA temp_store = MEMORY",
		"PRAGMA mmap_size = 268435456", // 256MB mmap
	}
	
	for _, pragma := range pragmas {
		_, err = suite.db.Exec(pragma)
		require.NoError(suite.T(), err, "Failed to set pragma: %s", pragma)
	}
}

// TearDownTest cleans up after each test
func (suite *DatabaseTestSuite) TearDownTest() {
	if suite.db != nil {
		suite.db.Close()
		suite.db = nil
	}
	if suite.persistenceDB != nil {
		suite.persistenceDB = nil
	}
}

// TestMigrationUpAndDown tests database schema migrations
func (suite *DatabaseTestSuite) TestMigrationUpAndDown() {
	suite.Run("MigrateUp", func() {
		// Set migration directory
		migrationsDir := "../../internal/manager/persistence/migrations"
		goose.SetDialect("sqlite3")
		
		// Test migration up
		err := goose.Up(suite.db, migrationsDir)
		require.NoError(suite.T(), err, "Failed to migrate up")
		
		// Verify current version
		version, err := goose.GetDBVersion(suite.db)
		require.NoError(suite.T(), err)
		assert.Greater(suite.T(), version, int64(0), "Database version should be greater than 0")
		
		suite.T().Logf("Migrated to version: %d", version)
		
		// Verify key tables exist
		expectedTables := []string{
			"goose_db_version",
			"jobs",
			"workers",
			"tasks",
			"worker_tags",
			"job_blocks",
			"task_failures",
			"worker_clusters",
			"sleep_schedules",
		}
		
		for _, tableName := range expectedTables {
			exists := suite.tableExists(tableName)
			assert.True(suite.T(), exists, "Table %s should exist after migration", tableName)
		}
	})
	
	suite.Run("MigrateDown", func() {
		// First migrate up to latest
		migrationsDir := "../../internal/manager/persistence/migrations"
		goose.SetDialect("sqlite3")
		
		err := goose.Up(suite.db, migrationsDir)
		require.NoError(suite.T(), err)
		
		initialVersion, err := goose.GetDBVersion(suite.db)
		require.NoError(suite.T(), err)
		
		// Test migration down (one step)
		err = goose.Down(suite.db, migrationsDir)
		require.NoError(suite.T(), err, "Failed to migrate down")
		
		// Verify version decreased
		newVersion, err := goose.GetDBVersion(suite.db)
		require.NoError(suite.T(), err)
		assert.Less(suite.T(), newVersion, initialVersion, "Version should decrease after down migration")
		
		suite.T().Logf("Migrated down from %d to %d", initialVersion, newVersion)
	})
	
	suite.Run("MigrationIdempotency", func() {
		migrationsDir := "../../internal/manager/persistence/migrations"
		goose.SetDialect("sqlite3")
		
		// Migrate up twice - should be safe
		err := goose.Up(suite.db, migrationsDir)
		require.NoError(suite.T(), err)
		
		version1, err := goose.GetDBVersion(suite.db)
		require.NoError(suite.T(), err)
		
		// Second migration up should not change anything
		err = goose.Up(suite.db, migrationsDir)
		require.NoError(suite.T(), err)
		
		version2, err := goose.GetDBVersion(suite.db)
		require.NoError(suite.T(), err)
		
		assert.Equal(suite.T(), version1, version2, "Multiple up migrations should be idempotent")
	})
}

// TestDataIntegrity tests data consistency and constraints
func (suite *DatabaseTestSuite) TestDataIntegrity() {
	// Migrate database first
	suite.migrateDatabase()
	
	// Initialize persistence layer
	var err error
	suite.persistenceDB, err = persistence.OpenDB(context.Background(), suite.testDBPath)
	require.NoError(suite.T(), err)
	
	suite.Run("ForeignKeyConstraints", func() {
		// Test foreign key relationships
		suite.testForeignKeyConstraints()
	})
	
	suite.Run("UniqueConstraints", func() {
		// Test unique constraints
		suite.testUniqueConstraints()
	})
	
	suite.Run("DataConsistency", func() {
		// Test data consistency across operations
		suite.testDataConsistency()
	})
	
	suite.Run("TransactionIntegrity", func() {
		// Test transaction rollback scenarios
		suite.testTransactionIntegrity()
	})
}

// TestConcurrentOperations tests database behavior under concurrent load
func (suite *DatabaseTestSuite) TestConcurrentOperations() {
	suite.migrateDatabase()
	
	var err error
	suite.persistenceDB, err = persistence.OpenDB(context.Background(), suite.testDBPath)
	require.NoError(suite.T(), err)
	
	suite.Run("ConcurrentJobCreation", func() {
		const numJobs = 50
		const concurrency = 10
		
		results := make(chan error, numJobs)
		sem := make(chan struct{}, concurrency)
		
		for i := 0; i < numJobs; i++ {
			go func(jobIndex int) {
				sem <- struct{}{}
				defer func() { <-sem }()
				
				ctx := context.Background()
				job := api.SubmittedJob{
					Name:              fmt.Sprintf("Concurrent Job %d", jobIndex),
					Type:              "test-job",
					Priority:          50,
					SubmitterPlatform: "linux",
					Settings:          map[string]interface{}{"test": "value"},
				}
				
				_, err := suite.persistenceDB.StoreJob(ctx, job)
				results <- err
			}(i)
		}
		
		// Collect results
		for i := 0; i < numJobs; i++ {
			err := <-results
			assert.NoError(suite.T(), err, "Concurrent job creation should succeed")
		}
		
		// Verify all jobs were created
		jobs, err := suite.persistenceDB.QueryJobs(context.Background(), api.JobsQuery{})
		require.NoError(suite.T(), err)
		assert.Len(suite.T(), jobs.Jobs, numJobs)
	})
	
	suite.Run("ConcurrentTaskUpdates", func() {
		// Create a job first
		ctx := context.Background()
		job := api.SubmittedJob{
			Name:              "Task Update Test Job",
			Type:              "test-job",
			Priority:          50,
			SubmitterPlatform: "linux",
			Settings:          map[string]interface{}{"frames": "1-10"},
		}
		
		storedJob, err := suite.persistenceDB.StoreJob(ctx, job)
		require.NoError(suite.T(), err)
		
		// Get tasks for the job
		tasks, err := suite.persistenceDB.QueryTasksByJobID(ctx, storedJob.Id)
		require.NoError(suite.T(), err)
		require.Greater(suite.T(), len(tasks), 0, "Should have tasks")
		
		// Concurrent task updates
		const numUpdates = 20
		results := make(chan error, numUpdates)
		
		for i := 0; i < numUpdates; i++ {
			go func(updateIndex int) {
				taskUpdate := api.TaskUpdate{
					TaskStatus: api.TaskStatusActive,
					Log:        fmt.Sprintf("Update %d", updateIndex),
					TaskProgress: &api.TaskProgress{
						PercentageComplete: int32(updateIndex * 5),
					},
				}
				
				err := suite.persistenceDB.UpdateTask(ctx, tasks[0].Uuid, taskUpdate)
				results <- err
			}(i)
		}
		
		// Collect results
		for i := 0; i < numUpdates; i++ {
			err := <-results
			assert.NoError(suite.T(), err, "Concurrent task updates should succeed")
		}
	})
}

// TestDatabasePerformance tests query performance and optimization
func (suite *DatabaseTestSuite) TestDatabasePerformance() {
	suite.migrateDatabase()
	
	var err error
	suite.persistenceDB, err = persistence.OpenDB(context.Background(), suite.testDBPath)
	require.NoError(suite.T(), err)
	
	suite.Run("QueryPerformance", func() {
		// Create test data
		suite.createTestData(100, 10, 500) // 100 jobs, 10 workers, 500 tasks
		
		// Test query performance
		performanceTests := []struct {
			name     string
			testFunc func() error
			maxTime  time.Duration
		}{
			{
				name: "QueryJobs",
				testFunc: func() error {
					_, err := suite.persistenceDB.QueryJobs(context.Background(), api.JobsQuery{})
					return err
				},
				maxTime: 100 * time.Millisecond,
			},
			{
				name: "QueryWorkers",
				testFunc: func() error {
					_, err := suite.persistenceDB.QueryWorkers(context.Background())
					return err
				},
				maxTime: 50 * time.Millisecond,
			},
			{
				name: "JobTasksSummary",
				testFunc: func() error {
					jobs, err := suite.persistenceDB.QueryJobs(context.Background(), api.JobsQuery{})
					if err != nil || len(jobs.Jobs) == 0 {
						return err
					}
					_, err = suite.persistenceDB.TaskStatsSummaryForJob(context.Background(), jobs.Jobs[0].Id)
					return err
				},
				maxTime: 50 * time.Millisecond,
			},
		}
		
		for _, test := range performanceTests {
			suite.T().Run(test.name, func(t *testing.T) {
				startTime := time.Now()
				err := test.testFunc()
				duration := time.Since(startTime)
				
				assert.NoError(t, err, "Query should succeed")
				assert.Less(t, duration, test.maxTime, 
					"Query %s took %v, should be under %v", test.name, duration, test.maxTime)
				
				t.Logf("Query %s completed in %v", test.name, duration)
			})
		}
	})
	
	suite.Run("IndexEfficiency", func() {
		// Test that indexes are being used effectively
		suite.analyzeQueryPlans()
	})
}

// TestDatabaseBackupRestore tests backup and restore functionality
func (suite *DatabaseTestSuite) TestDatabaseBackupRestore() {
	suite.migrateDatabase()
	
	var err error
	suite.persistenceDB, err = persistence.OpenDB(context.Background(), suite.testDBPath)
	require.NoError(suite.T(), err)
	
	suite.Run("BackupAndRestore", func() {
		// Create test data
		ctx := context.Background()
		originalJob := api.SubmittedJob{
			Name:              "Backup Test Job",
			Type:              "test-job",
			Priority:          50,
			SubmitterPlatform: "linux",
			Settings:          map[string]interface{}{"test": "backup"},
		}
		
		storedJob, err := suite.persistenceDB.StoreJob(ctx, originalJob)
		require.NoError(suite.T(), err)
		
		// Create backup
		backupPath := filepath.Join(suite.testHelper.TempDir(), "backup.sqlite")
		err = suite.createDatabaseBackup(suite.testDBPath, backupPath)
		require.NoError(suite.T(), err)
		
		// Verify backup exists and has data
		assert.FileExists(suite.T(), backupPath)
		
		// Test restore by opening backup database
		backupDB, err := sql.Open("sqlite", backupPath)
		require.NoError(suite.T(), err)
		defer backupDB.Close()
		
		// Verify data exists in backup
		var count int
		err = backupDB.QueryRow("SELECT COUNT(*) FROM jobs WHERE uuid = ?", storedJob.Id).Scan(&count)
		require.NoError(suite.T(), err)
		assert.Equal(suite.T(), 1, count, "Backup should contain the test job")
	})
}

// Helper methods

func (suite *DatabaseTestSuite) migrateDatabase() {
	migrationsDir := "../../internal/manager/persistence/migrations"
	goose.SetDialect("sqlite3")
	
	err := goose.Up(suite.db, migrationsDir)
	require.NoError(suite.T(), err, "Failed to migrate database")
}

func (suite *DatabaseTestSuite) tableExists(tableName string) bool {
	var count int
	query := `SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?`
	err := suite.db.QueryRow(query, tableName).Scan(&count)
	return err == nil && count > 0
}

func (suite *DatabaseTestSuite) testForeignKeyConstraints() {
	ctx := context.Background()
	
	// Test job-task relationship
	job := api.SubmittedJob{
		Name:              "FK Test Job",
		Type:              "test-job",
		Priority:          50,
		SubmitterPlatform: "linux",
		Settings:          map[string]interface{}{"test": "fk"},
	}
	
	storedJob, err := suite.persistenceDB.StoreJob(ctx, job)
	require.NoError(suite.T(), err)
	
	// Delete job should handle tasks appropriately
	err = suite.persistenceDB.DeleteJob(ctx, storedJob.Id)
	require.NoError(suite.T(), err)
	
	// Verify job and related tasks are handled correctly
	_, err = suite.persistenceDB.FetchJob(ctx, storedJob.Id)
	assert.Error(suite.T(), err, "Job should not exist after deletion")
}

func (suite *DatabaseTestSuite) testUniqueConstraints() {
	ctx := context.Background()
	
	// Test duplicate job names (should be allowed)
	job1 := api.SubmittedJob{
		Name:              "Duplicate Name Test",
		Type:              "test-job",
		Priority:          50,
		SubmitterPlatform: "linux",
		Settings:          map[string]interface{}{"test": "unique1"},
	}
	
	job2 := api.SubmittedJob{
		Name:              "Duplicate Name Test", // Same name
		Type:              "test-job",
		Priority:          50,
		SubmitterPlatform: "linux",
		Settings:          map[string]interface{}{"test": "unique2"},
	}
	
	_, err1 := suite.persistenceDB.StoreJob(ctx, job1)
	_, err2 := suite.persistenceDB.StoreJob(ctx, job2)
	
	assert.NoError(suite.T(), err1, "First job should be stored successfully")
	assert.NoError(suite.T(), err2, "Duplicate job names should be allowed")
}

func (suite *DatabaseTestSuite) testDataConsistency() {
	ctx := context.Background()
	
	// Create job with tasks
	job := api.SubmittedJob{
		Name:              "Consistency Test Job",
		Type:              "simple-blender-render",
		Priority:          50,
		SubmitterPlatform: "linux",
		Settings: map[string]interface{}{
			"frames": "1-5",
			"chunk_size": 1,
		},
	}
	
	storedJob, err := suite.persistenceDB.StoreJob(ctx, job)
	require.NoError(suite.T(), err)
	
	// Verify tasks were created
	tasks, err := suite.persistenceDB.QueryTasksByJobID(ctx, storedJob.Id)
	require.NoError(suite.T(), err)
	assert.Greater(suite.T(), len(tasks), 0, "Job should have tasks")
	
	// Update task status and verify job status reflects changes
	if len(tasks) > 0 {
		taskUpdate := api.TaskUpdate{
			TaskStatus: api.TaskStatusCompleted,
			Log:        "Task completed for consistency test",
		}
		
		err = suite.persistenceDB.UpdateTask(ctx, tasks[0].Uuid, taskUpdate)
		require.NoError(suite.T(), err)
		
		// Check job status was updated appropriately
		updatedJob, err := suite.persistenceDB.FetchJob(ctx, storedJob.Id)
		require.NoError(suite.T(), err)
		
		// Job status should reflect task progress
		assert.NotEqual(suite.T(), api.JobStatusQueued, updatedJob.Status, 
			"Job status should change when tasks are updated")
	}
}

func (suite *DatabaseTestSuite) testTransactionIntegrity() {
	ctx := context.Background()
	
	// Test transaction rollback on constraint violation
	tx, err := suite.db.BeginTx(ctx, nil)
	require.NoError(suite.T(), err)
	
	// Insert valid data
	_, err = tx.Exec("INSERT INTO jobs (uuid, name, job_type, priority, status, created) VALUES (?, ?, ?, ?, ?, ?)",
		"test-tx-1", "Transaction Test", "test", 50, "queued", time.Now().UTC())
	require.NoError(suite.T(), err)
	
	// Attempt to insert invalid data (this should cause rollback)
	_, err = tx.Exec("INSERT INTO tasks (uuid, job_id, name, task_type, status) VALUES (?, ?, ?, ?, ?)",
		"test-task-1", "non-existent-job", "Test Task", "test", "queued")
	
	if err != nil {
		// Rollback transaction
		tx.Rollback()
		
		// Verify original data was not committed
		var count int
		suite.db.QueryRow("SELECT COUNT(*) FROM jobs WHERE uuid = ?", "test-tx-1").Scan(&count)
		assert.Equal(suite.T(), 0, count, "Transaction should be rolled back on constraint violation")
	} else {
		tx.Commit()
	}
}

func (suite *DatabaseTestSuite) createTestData(numJobs, numWorkers, numTasks int) {
	ctx := context.Background()
	
	// Create jobs
	for i := 0; i < numJobs; i++ {
		job := api.SubmittedJob{
			Name:              fmt.Sprintf("Performance Test Job %d", i),
			Type:              "test-job",
			Priority:          50,
			SubmitterPlatform: "linux",
			Settings:          map[string]interface{}{"frames": "1-10"},
		}
		
		_, err := suite.persistenceDB.StoreJob(ctx, job)
		require.NoError(suite.T(), err)
	}
	
	suite.T().Logf("Created %d test jobs", numJobs)
}

func (suite *DatabaseTestSuite) analyzeQueryPlans() {
	// Common queries to analyze
	queries := []string{
		"SELECT * FROM jobs WHERE status = 'queued'",
		"SELECT * FROM tasks WHERE job_id = 'some-job-id'",
		"SELECT * FROM workers WHERE status = 'awake'",
		"SELECT job_id, COUNT(*) FROM tasks GROUP BY job_id",
	}
	
	for _, query := range queries {
		explainQuery := "EXPLAIN QUERY PLAN " + query
		rows, err := suite.db.Query(explainQuery)
		if err != nil {
			suite.T().Logf("Failed to explain query: %s, error: %v", query, err)
			continue
		}
		
		suite.T().Logf("Query plan for: %s", query)
		for rows.Next() {
			var id, parent, notused int
			var detail string
			rows.Scan(&id, &parent, &notused, &detail)
			suite.T().Logf("  %s", detail)
		}
		rows.Close()
	}
}

func (suite *DatabaseTestSuite) createDatabaseBackup(sourcePath, backupPath string) error {
	// Simple file copy for SQLite
	sourceFile, err := os.Open(sourcePath)
	if err != nil {
		return err
	}
	defer sourceFile.Close()
	
	backupFile, err := os.Create(backupPath)
	if err != nil {
		return err
	}
	defer backupFile.Close()
	
	_, err = backupFile.ReadFrom(sourceFile)
	return err
}

// TestLargeDataOperations tests database behavior with large datasets
func (suite *DatabaseTestSuite) TestLargeDataOperations() {
	suite.migrateDatabase()
	
	var err error
	suite.persistenceDB, err = persistence.OpenDB(context.Background(), suite.testDBPath)
	require.NoError(suite.T(), err)
	
	suite.Run("LargeJobWithManyTasks", func() {
		ctx := context.Background()
		
		// Create job with many frames
		job := api.SubmittedJob{
			Name:              "Large Frame Job",
			Type:              "simple-blender-render",
			Priority:          50,
			SubmitterPlatform: "linux",
			Settings: map[string]interface{}{
				"frames":     "1-1000", // 1000 frames
				"chunk_size": 10,       // 100 tasks
			},
		}
		
		startTime := time.Now()
		storedJob, err := suite.persistenceDB.StoreJob(ctx, job)
		creationTime := time.Since(startTime)
		
		require.NoError(suite.T(), err)
		assert.Less(suite.T(), creationTime, 5*time.Second, 
			"Large job creation should complete within 5 seconds")
		
		// Verify tasks were created
		tasks, err := suite.persistenceDB.QueryTasksByJobID(ctx, storedJob.Id)
		require.NoError(suite.T(), err)
		assert.Greater(suite.T(), len(tasks), 90, "Should create around 100 tasks")
		
		suite.T().Logf("Created job with %d tasks in %v", len(tasks), creationTime)
	})
	
	suite.Run("BulkTaskUpdates", func() {
		// Test updating many tasks efficiently
		ctx := context.Background()
		
		job := api.SubmittedJob{
			Name:              "Bulk Update Test Job",
			Type:              "simple-blender-render",
			Priority:          50,
			SubmitterPlatform: "linux",
			Settings: map[string]interface{}{
				"frames":     "1-100",
				"chunk_size": 5,
			},
		}
		
		storedJob, err := suite.persistenceDB.StoreJob(ctx, job)
		require.NoError(suite.T(), err)
		
		tasks, err := suite.persistenceDB.QueryTasksByJobID(ctx, storedJob.Id)
		require.NoError(suite.T(), err)
		
		// Update all tasks
		startTime := time.Now()
		for _, task := range tasks {
			taskUpdate := api.TaskUpdate{
				TaskStatus: api.TaskStatusCompleted,
				Log:        "Bulk update test completed",
			}
			
			err := suite.persistenceDB.UpdateTask(ctx, task.Uuid, taskUpdate)
			require.NoError(suite.T(), err)
		}
		updateTime := time.Since(startTime)
		
		assert.Less(suite.T(), updateTime, 2*time.Second, 
			"Bulk task updates should complete efficiently")
		
		suite.T().Logf("Updated %d tasks in %v", len(tasks), updateTime)
	})
}

// TestSuite runs all database tests
func TestDatabaseSuite(t *testing.T) {
	suite.Run(t, new(DatabaseTestSuite))
}