c74e91dd69
Test improvements: - fixtures/: Updated mocks, fixtures with group context, SSH server, TUI driver - integration/: WebSocket queue and handler tests with groups - e2e/: WebSocket and TLS proxy end-to-end tests - unit/api/ws_test.go: WebSocket API tests - unit/scheduler/service_templates_test.go: Service template tests - benchmarks/scheduler_bench_test.go: Performance benchmarks Cleanup: - Remove privacy middleware (replaced by audit system) - Remove privacy_test.go
128 lines
3.8 KiB
Go
128 lines
3.8 KiB
Go
// Package fixtures provides shared test utilities and fixtures for scheduler tests
|
|
package tests
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/scheduler"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// SchedulerTestFixture provides a test fixture for scheduler tests
|
|
type SchedulerTestFixture struct {
|
|
T testing.TB
|
|
Hub *scheduler.SchedulerHub
|
|
Workers map[string]*MockWorker
|
|
stateDir string
|
|
}
|
|
|
|
// NewSchedulerTestFixture creates a new scheduler test fixture
|
|
func NewSchedulerTestFixture(t testing.TB, cfg scheduler.HubConfig) *SchedulerTestFixture {
|
|
if cfg.BindAddr == "" {
|
|
cfg.BindAddr = "localhost:0"
|
|
}
|
|
|
|
// Create isolated state directory per test
|
|
stateDir, err := os.MkdirTemp("", "fetchml-test-*")
|
|
require.NoError(t, err)
|
|
cfg.StateDir = stateDir
|
|
|
|
hub, err := scheduler.NewHub(cfg, nil)
|
|
require.NoError(t, err)
|
|
|
|
// Start scheduler
|
|
err = hub.Start()
|
|
require.NoError(t, err)
|
|
|
|
return &SchedulerTestFixture{
|
|
T: t,
|
|
Hub: hub,
|
|
Workers: make(map[string]*MockWorker),
|
|
stateDir: stateDir,
|
|
}
|
|
}
|
|
|
|
// CreateWorker creates and registers a new mock worker
|
|
func (f *SchedulerTestFixture) CreateWorker(workerID string, caps scheduler.WorkerCapabilities) *MockWorker {
|
|
worker := NewMockWorker(f.T, f.Hub, workerID)
|
|
worker.Register(caps)
|
|
f.Workers[workerID] = worker
|
|
return worker
|
|
}
|
|
|
|
// SubmitJob submits a job to the scheduler
|
|
func (f *SchedulerTestFixture) SubmitJob(spec scheduler.JobSpec) {
|
|
err := f.Hub.SubmitJob(spec)
|
|
require.NoError(f.T, err)
|
|
}
|
|
|
|
// GetTask retrieves a task by ID
|
|
func (f *SchedulerTestFixture) GetTask(taskID string) *scheduler.Task {
|
|
return f.Hub.GetTask(taskID)
|
|
}
|
|
|
|
// Cleanup stops the scheduler, closes all workers, and removes state dir
|
|
func (f *SchedulerTestFixture) Cleanup() {
|
|
// Close all workers first
|
|
for _, worker := range f.Workers {
|
|
worker.Close()
|
|
}
|
|
// Then stop the hub
|
|
f.Hub.Stop()
|
|
// Clean up isolated state directory
|
|
if err := os.RemoveAll(f.stateDir); err != nil {
|
|
// Log cleanup error but don't fail test
|
|
fmt.Fprintf(os.Stderr, "failed to remove state dir: %v\n", err)
|
|
}
|
|
}
|
|
|
|
// DefaultHubConfig returns a default hub configuration for testing
|
|
func DefaultHubConfig() scheduler.HubConfig {
|
|
tokens := map[string]string{
|
|
"test-token-worker-restart-1": "worker-restart-1",
|
|
"test-token-mode-switch-worker": "mode-switch-worker",
|
|
"test-token-mode-switch-worker-2": "mode-switch-worker-2",
|
|
"test-token-e2e-worker-1": "e2e-worker-1",
|
|
"test-token-e2e-worker-2": "e2e-worker-2",
|
|
"test-token-worker-death-test": "worker-death-test",
|
|
"test-token-worker-split-1": "worker-split-1",
|
|
"test-token-worker-split-2": "worker-split-2",
|
|
"test-token-worker-split-3": "worker-split-3",
|
|
"test-token-worker-timeout": "worker-timeout",
|
|
"test-token-worker-gang": "worker-gang",
|
|
"test-token-bench-worker": "bench-worker",
|
|
"test-token-bench-hb-worker": "bench-hb-worker",
|
|
"test-token-bench-assign-worker": "bench-assign-worker",
|
|
}
|
|
|
|
// Add tokens for dynamic benchmark worker IDs (0-999 for each pattern)
|
|
for i := range 1000 {
|
|
tokens[fmt.Sprintf("test-token-bench-worker-%d", i)] = fmt.Sprintf("bench-worker-%d", i)
|
|
tokens[fmt.Sprintf("test-token-bench-multi-worker-%d", i)] = fmt.Sprintf("bench-multi-worker-%d", i)
|
|
}
|
|
|
|
return scheduler.HubConfig{
|
|
BindAddr: "localhost:0",
|
|
DefaultBatchSlots: 4,
|
|
StarvationThresholdMins: 5,
|
|
AcceptanceTimeoutSecs: 5,
|
|
GangAllocTimeoutSecs: 10,
|
|
// #nosec G101 -- These are test fixture tokens, not real credentials
|
|
WorkerTokens: tokens,
|
|
}
|
|
}
|
|
|
|
// WaitForTimeout is a helper to wait for a condition with timeout
|
|
func WaitForTimeout(duration time.Duration, condition func() bool) bool {
|
|
deadline := time.Now().Add(duration)
|
|
for time.Now().Before(deadline) {
|
|
if condition() {
|
|
return true
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
return false
|
|
}
|