flamenco/tests/helpers/test_helper.go

package helpers

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

import (
	"context"
	"database/sql"
	"fmt"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/pressly/goose/v3"
	_ "modernc.org/sqlite"

	"projects.blender.org/studio/flamenco/internal/manager/api_impl"
	"projects.blender.org/studio/flamenco/internal/manager/config"
	"projects.blender.org/studio/flamenco/internal/manager/job_compilers"
	"projects.blender.org/studio/flamenco/internal/manager/persistence"
	"projects.blender.org/studio/flamenco/pkg/api"
)

// TestHelper provides common testing utilities and setup
type TestHelper struct {
	t         *testing.T
	tempDir   string
	server    *httptest.Server
	dbPath    string
	db        *persistence.DB
	cleanup   []func()
}

// TestFixtures contains common test data
type TestFixtures struct {
	Jobs    []api.SubmittedJob
	Workers []api.WorkerRegistration
	Tasks   []api.Task
}

// NewTestHelper creates a new test helper instance
func NewTestHelper(t *testing.T) *TestHelper {
	helper := &TestHelper{
		t:       t,
		cleanup: make([]func(), 0),
	}
	
	// Create temporary directory for test files
	helper.createTempDir()
	
	return helper
}

// CreateTempDir creates a temporary directory for tests
func (h *TestHelper) CreateTempDir(suffix string) string {
	if h.tempDir == "" {
		h.createTempDir()
	}
	
	dir := filepath.Join(h.tempDir, suffix)
	err := os.MkdirAll(dir, 0755)
	if err != nil {
		h.t.Fatalf("Failed to create temp directory: %v", err)
	}
	
	return dir
}

// TempDir returns the temporary directory path
func (h *TestHelper) TempDir() string {
	return h.tempDir
}

// StartTestServer starts a test HTTP server with Flamenco Manager
func (h *TestHelper) StartTestServer() *httptest.Server {
	if h.server != nil {
		return h.server
	}
	
	// Setup test database
	h.setupTestDatabase()
	
	// Create test configuration
	cfg := h.createTestConfig()
	
	// Setup Echo server
	e := echo.New()
	e.HideBanner = true
	
	// Setup API implementation with test dependencies
	flamenco := h.createTestFlamenco(cfg)
	api.RegisterHandlers(e, flamenco)
	
	// Start test server
	h.server = httptest.NewServer(e)
	h.addCleanup(func() {
		h.server.Close()
		h.server = nil
	})
	
	return h.server
}

// SetupTestDatabase creates and migrates a test database
func (h *TestHelper) setupTestDatabase() *persistence.DB {
	if h.db != nil {
		return h.db
	}
	
	// Create test database path
	h.dbPath = filepath.Join(h.tempDir, "test_flamenco.sqlite")
	
	// Remove existing database
	os.Remove(h.dbPath)
	
	// Open database connection
	sqlDB, err := sql.Open("sqlite", h.dbPath)
	if err != nil {
		h.t.Fatalf("Failed to open test database: %v", err)
	}
	
	// Set SQLite pragmas for testing
	pragmas := []string{
		"PRAGMA foreign_keys = ON",
		"PRAGMA journal_mode = WAL",
		"PRAGMA synchronous = NORMAL",
		"PRAGMA cache_size = -32000", // 32MB cache
	}
	
	for _, pragma := range pragmas {
		_, err = sqlDB.Exec(pragma)
		if err != nil {
			h.t.Fatalf("Failed to set pragma %s: %v", pragma, err)
		}
	}
	
	// Run migrations
	migrationsDir := h.findMigrationsDir()
	goose.SetDialect("sqlite3")
	
	err = goose.Up(sqlDB, migrationsDir)
	if err != nil {
		h.t.Fatalf("Failed to migrate test database: %v", err)
	}
	
	sqlDB.Close()
	
	// Open with persistence layer
	h.db, err = persistence.OpenDB(context.Background(), h.dbPath)
	if err != nil {
		h.t.Fatalf("Failed to open persistence DB: %v", err)
	}
	
	h.addCleanup(func() {
		if h.db != nil {
			h.db.Close()
			h.db = nil
		}
	})
	
	return h.db
}

// GetTestDatabase returns the test database instance
func (h *TestHelper) GetTestDatabase() *persistence.DB {
	if h.db == nil {
		h.setupTestDatabase()
	}
	return h.db
}

// CreateTestJob creates a test job with reasonable defaults
func (h *TestHelper) CreateTestJob(name string, jobType string) api.SubmittedJob {
	return api.SubmittedJob{
		Name:              name,
		Type:              jobType,
		Priority:          50,
		SubmitterPlatform: "linux",
		Settings: map[string]interface{}{
			"filepath":               "/shared-storage/test.blend",
			"chunk_size":             5,
			"format":                 "PNG",
			"image_file_extension":   ".png",
			"frames":                 "1-10",
			"render_output_root":     "/shared-storage/renders/",
			"add_path_components":    0,
			"render_output_path":     "/shared-storage/renders/test/######",
		},
	}
}

// CreateTestWorker creates a test worker registration
func (h *TestHelper) CreateTestWorker(name string) api.WorkerRegistration {
	return api.WorkerRegistration{
		Name:               name,
		Address:            "192.168.1.100",
		Platform:           "linux",
		SoftwareVersion:    "3.0.0",
		SupportedTaskTypes: []string{"blender", "ffmpeg", "file-management"},
	}
}

// LoadTestFixtures loads common test data fixtures
func (h *TestHelper) LoadTestFixtures() *TestFixtures {
	return &TestFixtures{
		Jobs: []api.SubmittedJob{
			h.CreateTestJob("Test Animation Render", "simple-blender-render"),
			h.CreateTestJob("Test Still Render", "simple-blender-render"),
			h.CreateTestJob("Test Video Encode", "simple-blender-render"),
		},
		Workers: []api.WorkerRegistration{
			h.CreateTestWorker("test-worker-1"),
			h.CreateTestWorker("test-worker-2"),
			h.CreateTestWorker("test-worker-3"),
		},
	}
}

// WaitForCondition waits for a condition to become true with timeout
func (h *TestHelper) WaitForCondition(timeout time.Duration, condition func() bool) bool {
	deadline := time.After(timeout)
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()
	
	for {
		select {
		case <-deadline:
			return false
		case <-ticker.C:
			if condition() {
				return true
			}
		}
	}
}

// AssertEventuallyTrue waits for a condition and fails test if not met
func (h *TestHelper) AssertEventuallyTrue(timeout time.Duration, condition func() bool, message string) {
	if !h.WaitForCondition(timeout, condition) {
		h.t.Fatalf("Condition not met within %v: %s", timeout, message)
	}
}

// CreateTestFiles creates test files in the temporary directory
func (h *TestHelper) CreateTestFiles(files map[string]string) {
	for filename, content := range files {
		fullPath := filepath.Join(h.tempDir, filename)
		
		// Create directory if needed
		dir := filepath.Dir(fullPath)
		err := os.MkdirAll(dir, 0755)
		if err != nil {
			h.t.Fatalf("Failed to create directory %s: %v", dir, err)
		}
		
		// Write file
		err = os.WriteFile(fullPath, []byte(content), 0644)
		if err != nil {
			h.t.Fatalf("Failed to create test file %s: %v", fullPath, err)
		}
	}
}

// Cleanup runs all registered cleanup functions
func (h *TestHelper) Cleanup() {
	for i := len(h.cleanup) - 1; i >= 0; i-- {
		h.cleanup[i]()
	}
	
	// Remove temporary directory
	if h.tempDir != "" {
		os.RemoveAll(h.tempDir)
		h.tempDir = ""
	}
}

// Private helper methods

func (h *TestHelper) createTempDir() {
	var err error
	h.tempDir, err = os.MkdirTemp("", "flamenco-test-*")
	if err != nil {
		h.t.Fatalf("Failed to create temp directory: %v", err)
	}
}

func (h *TestHelper) addCleanup(fn func()) {
	h.cleanup = append(h.cleanup, fn)
}

func (h *TestHelper) findMigrationsDir() string {
	// Try different relative paths to find migrations
	candidates := []string{
		"../../internal/manager/persistence/migrations",
		"../internal/manager/persistence/migrations", 
		"./internal/manager/persistence/migrations",
		"internal/manager/persistence/migrations",
	}
	
	for _, candidate := range candidates {
		if _, err := os.Stat(candidate); err == nil {
			return candidate
		}
	}
	
	h.t.Fatalf("Could not find migrations directory")
	return ""
}

func (h *TestHelper) createTestConfig() *config.Conf {
	cfg := &config.Conf{
		Base: config.Base{
			DatabaseDSN:      h.dbPath,
			SharedStoragePath: filepath.Join(h.tempDir, "shared-storage"),
		},
		Manager: config.Manager{
			DatabaseCheckPeriod: config.Duration{Duration: 1 * time.Minute},
		},
	}
	
	// Create shared storage directory
	err := os.MkdirAll(cfg.Base.SharedStoragePath, 0755)
	if err != nil {
		h.t.Fatalf("Failed to create shared storage directory: %v", err)
	}
	
	return cfg
}

func (h *TestHelper) createTestFlamenco(cfg *config.Conf) api_impl.ServerInterface {
	// This is a simplified test setup
	// In a real implementation, you'd wire up all dependencies properly
	flamenco := &TestFlamencoImpl{
		config:   cfg,
		database: h.GetTestDatabase(),
	}
	
	return flamenco
}

// TestFlamencoImpl provides a minimal implementation for testing
type TestFlamencoImpl struct {
	config   *config.Conf
	database *persistence.DB
}

// Implement minimal ServerInterface methods for testing
func (f *TestFlamencoImpl) GetVersion(ctx echo.Context) error {
	version := api.FlamencoVersion{
		Version: "3.0.0-test",
		Name:    "flamenco",
	}
	return ctx.JSON(200, version)
}

func (f *TestFlamencoImpl) GetConfiguration(ctx echo.Context) error {
	cfg := api.ManagerConfiguration{
		Variables: map[string]api.ManagerVariable{
			"blender": {
				IsTwoWay: false,
				Values: []api.ManagerVariableValue{
					{
						Platform: "linux",
						Value:    "/usr/local/blender/blender",
					},
				},
			},
		},
	}
	return ctx.JSON(200, cfg)
}

func (f *TestFlamencoImpl) SubmitJob(ctx echo.Context) error {
	var submittedJob api.SubmittedJob
	if err := ctx.Bind(&submittedJob); err != nil {
		return ctx.JSON(400, map[string]string{"error": "Invalid job data"})
	}
	
	// Store job in database
	job, err := f.database.StoreJob(ctx.Request().Context(), submittedJob)
	if err != nil {
		return ctx.JSON(500, map[string]string{"error": "Failed to store job"})
	}
	
	return ctx.JSON(200, job)
}

func (f *TestFlamencoImpl) QueryJobs(ctx echo.Context) error {
	jobs, err := f.database.QueryJobs(ctx.Request().Context(), api.JobsQuery{})
	if err != nil {
		return ctx.JSON(500, map[string]string{"error": "Failed to query jobs"})
	}
	
	return ctx.JSON(200, jobs)
}

func (f *TestFlamencoImpl) FetchJob(ctx echo.Context, jobID string) error {
	job, err := f.database.FetchJob(ctx.Request().Context(), jobID)
	if err != nil {
		return ctx.JSON(404, map[string]string{"error": "Job not found"})
	}
	
	return ctx.JSON(200, job)
}

func (f *TestFlamencoImpl) DeleteJob(ctx echo.Context, jobID string) error {
	err := f.database.DeleteJob(ctx.Request().Context(), jobID)
	if err != nil {
		return ctx.JSON(404, map[string]string{"error": "Job not found"})
	}
	
	return ctx.NoContent(204)
}

func (f *TestFlamencoImpl) RegisterWorker(ctx echo.Context) error {
	var workerReg api.WorkerRegistration
	if err := ctx.Bind(&workerReg); err != nil {
		return ctx.JSON(400, map[string]string{"error": "Invalid worker data"})
	}
	
	worker, err := f.database.CreateWorker(ctx.Request().Context(), workerReg)
	if err != nil {
		return ctx.JSON(500, map[string]string{"error": "Failed to register worker"})
	}
	
	return ctx.JSON(200, worker)
}

func (f *TestFlamencoImpl) QueryWorkers(ctx echo.Context) error {
	workers, err := f.database.QueryWorkers(ctx.Request().Context())
	if err != nil {
		return ctx.JSON(500, map[string]string{"error": "Failed to query workers"})
	}
	
	return ctx.JSON(200, api.WorkerList{Workers: workers})
}

func (f *TestFlamencoImpl) SignOnWorker(ctx echo.Context, workerID string) error {
	var signOn api.WorkerSignOn
	if err := ctx.Bind(&signOn); err != nil {
		return ctx.JSON(400, map[string]string{"error": "Invalid sign-on data"})
	}
	
	// Simple sign-on implementation
	response := api.WorkerStateChange{
		StatusRequested: &[]api.WorkerStatus{api.WorkerStatusAwake}[0],
	}
	
	return ctx.JSON(200, response)
}

func (f *TestFlamencoImpl) ScheduleTask(ctx echo.Context, workerID string) error {
	// Simple task scheduling - return no content if no tasks available
	return ctx.NoContent(204)
}

func (f *TestFlamencoImpl) TaskUpdate(ctx echo.Context, workerID, taskID string) error {
	var taskUpdate api.TaskUpdate
	if err := ctx.Bind(&taskUpdate); err != nil {
		return ctx.JSON(400, map[string]string{"error": "Invalid task update"})
	}
	
	// Update task in database
	err := f.database.UpdateTask(ctx.Request().Context(), taskID, taskUpdate)
	if err != nil {
		return ctx.JSON(404, map[string]string{"error": "Task not found"})
	}
	
	return ctx.NoContent(204)
}

// Add other required methods as stubs
func (f *TestFlamencoImpl) CheckSharedStoragePath(ctx echo.Context) error {
	return ctx.JSON(200, map[string]interface{}{"is_usable": true})
}

func (f *TestFlamencoImpl) ShamanCheckout(ctx echo.Context) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) ShamanCheckoutRequirements(ctx echo.Context) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) ShamanFileStore(ctx echo.Context, checksum string, filesize int) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) ShamanFileStoreCheck(ctx echo.Context, checksum string, filesize int) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) GetJobType(ctx echo.Context, typeName string) error {
	// Return a simple job type for testing
	jobType := api.AvailableJobType{
		Name:  typeName,
		Label: fmt.Sprintf("Test %s", typeName),
		Settings: []api.AvailableJobSetting{
			{
				Key:      "filepath",
				Type:     api.AvailableJobSettingTypeString,
				Required: true,
			},
		},
	}
	return ctx.JSON(200, jobType)
}

func (f *TestFlamencoImpl) GetJobTypes(ctx echo.Context) error {
	jobTypes := api.AvailableJobTypes{
		JobTypes: []api.AvailableJobType{
			{
				Name:  "simple-blender-render",
				Label: "Simple Blender Render",
			},
			{
				Name:  "test-job",
				Label: "Test Job Type",
			},
		},
	}
	return ctx.JSON(200, jobTypes)
}

// Add placeholder methods for other required ServerInterface methods
func (f *TestFlamencoImpl) FetchTask(ctx echo.Context, taskID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) FetchTaskLogTail(ctx echo.Context, taskID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) FetchWorker(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) RequestWorkerStatusChange(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) DeleteWorker(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) FetchWorkerSleepSchedule(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) SetWorkerSleepSchedule(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) DeleteWorkerSleepSchedule(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) SetWorkerTags(ctx echo.Context, workerID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) GetVariables(ctx echo.Context, audience string, platform string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) QueryTasksByJobID(ctx echo.Context, jobID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) GetTaskLogInfo(ctx echo.Context, taskID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) FetchGlobalLastRenderedInfo(ctx echo.Context) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}

func (f *TestFlamencoImpl) FetchJobLastRenderedInfo(ctx echo.Context, jobID string) error {
	return ctx.JSON(501, map[string]string{"error": "Not implemented in test"})
}