fc2459977c
**Worker Refactoring:** - Update internal/worker/factory.go, worker.go, snapshot_store.go - Update native_bridge.go and native_bridge_nocgo.go for native library integration **Test Updates:** - Update all worker unit tests for new interfaces - Update chaos tests - Update container/podman_test.go - Add internal/workertest/worker.go for shared test utilities **Documentation:** - Update native/README.md
296 lines
8.6 KiB
Go
296 lines
8.6 KiB
Go
package worker_test
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/alicebob/miniredis/v2"
|
|
"github.com/jfraeys/fetch_ml/internal/queue"
|
|
"github.com/jfraeys/fetch_ml/internal/worker"
|
|
"github.com/jfraeys/fetch_ml/internal/workertest"
|
|
)
|
|
|
|
func TestPrewarmNextOnce_Snapshot_WritesPrewarmDir(t *testing.T) {
|
|
s, err := miniredis.Run()
|
|
if err != nil {
|
|
t.Fatalf("miniredis: %v", err)
|
|
}
|
|
t.Cleanup(s.Close)
|
|
|
|
base := t.TempDir()
|
|
dataDir := filepath.Join(base, "data")
|
|
|
|
// Create a snapshot directory and compute its overall SHA.
|
|
srcSnapshot := filepath.Join(base, "snapshot-src")
|
|
if err := os.MkdirAll(srcSnapshot, 0750); err != nil {
|
|
t.Fatalf("mkdir src snapshot: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(srcSnapshot, "file.txt"), []byte("ok"), 0600); err != nil {
|
|
t.Fatalf("write snapshot file: %v", err)
|
|
}
|
|
sha, err := worker.DirOverallSHA256Hex(srcSnapshot)
|
|
if err != nil {
|
|
t.Fatalf("DirOverallSHA256Hex: %v", err)
|
|
}
|
|
cacheDir := filepath.Join(dataDir, "snapshots", "sha256", sha)
|
|
if err := os.MkdirAll(filepath.Dir(cacheDir), 0750); err != nil {
|
|
t.Fatalf("mkdir cache parent: %v", err)
|
|
}
|
|
if err := os.Rename(srcSnapshot, cacheDir); err != nil {
|
|
t.Fatalf("rename into cache: %v", err)
|
|
}
|
|
|
|
// Queue has one task; prewarm should stage it into base/.prewarm/snapshots/<taskID>.
|
|
tq, err := queue.NewTaskQueue(queue.Config{RedisAddr: s.Addr(), MetricsFlushInterval: 5 * time.Millisecond})
|
|
if err != nil {
|
|
t.Fatalf("NewTaskQueue: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = tq.Close() })
|
|
|
|
task := &queue.Task{
|
|
ID: "task-1",
|
|
JobName: "job-1",
|
|
Status: "queued",
|
|
Priority: 10,
|
|
CreatedAt: time.Now().UTC(),
|
|
SnapshotID: "snap-1",
|
|
Metadata: map[string]string{
|
|
"snapshot_sha256": "sha256:" + sha,
|
|
},
|
|
}
|
|
if err := tq.AddTask(task); err != nil {
|
|
t.Fatalf("AddTask: %v", err)
|
|
}
|
|
|
|
cfg := &worker.Config{
|
|
WorkerID: "worker-1",
|
|
BasePath: base,
|
|
DataDir: dataDir,
|
|
PrewarmEnabled: true,
|
|
AutoFetchData: false,
|
|
LocalMode: true,
|
|
PollInterval: 1,
|
|
MaxWorkers: 1,
|
|
DatasetCacheTTL: 30 * time.Minute,
|
|
}
|
|
w := workertest.NewTestWorkerWithQueue(cfg, tq)
|
|
|
|
ok, err := w.PrewarmNextOnce(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("PrewarmNextOnce: %v", err)
|
|
}
|
|
if !ok {
|
|
t.Fatalf("expected ok=true")
|
|
}
|
|
|
|
prewarmed := filepath.Join(base, ".prewarm", "snapshots", task.ID, "file.txt")
|
|
if _, err := os.Stat(prewarmed); err != nil {
|
|
t.Fatalf("expected prewarmed file to exist: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestPrewarmNextOnce_Disabled_NoOp(t *testing.T) {
|
|
s, err := miniredis.Run()
|
|
if err != nil {
|
|
t.Fatalf("miniredis: %v", err)
|
|
}
|
|
t.Cleanup(s.Close)
|
|
|
|
base := t.TempDir()
|
|
dataDir := filepath.Join(base, "data")
|
|
|
|
tq, err := queue.NewTaskQueue(queue.Config{RedisAddr: s.Addr(), MetricsFlushInterval: 5 * time.Millisecond})
|
|
if err != nil {
|
|
t.Fatalf("NewTaskQueue: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = tq.Close() })
|
|
|
|
task := &queue.Task{ID: "task-1", JobName: "job-1", Status: "queued", Priority: 10, CreatedAt: time.Now().UTC()}
|
|
if err := tq.AddTask(task); err != nil {
|
|
t.Fatalf("AddTask: %v", err)
|
|
}
|
|
|
|
cfg := &worker.Config{WorkerID: "worker-1", BasePath: base, DataDir: dataDir, PrewarmEnabled: false}
|
|
w := workertest.NewTestWorkerWithQueue(cfg, tq)
|
|
|
|
ok, err := w.PrewarmNextOnce(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("PrewarmNextOnce: %v", err)
|
|
}
|
|
if ok {
|
|
t.Fatalf("expected ok=false")
|
|
}
|
|
if _, err := os.Stat(filepath.Join(base, ".prewarm")); err == nil {
|
|
t.Fatalf("expected no .prewarm dir when disabled")
|
|
}
|
|
}
|
|
|
|
func TestPrewarmNextOnce_QueueEmpty_DoesNotDeleteState(t *testing.T) {
|
|
s, err := miniredis.Run()
|
|
if err != nil {
|
|
t.Fatalf("miniredis: %v", err)
|
|
}
|
|
t.Cleanup(s.Close)
|
|
|
|
base := t.TempDir()
|
|
dataDir := filepath.Join(base, "data")
|
|
|
|
// Create a snapshot directory and compute its overall SHA.
|
|
srcSnapshot := filepath.Join(base, "snapshot-src")
|
|
if err := os.MkdirAll(srcSnapshot, 0750); err != nil {
|
|
t.Fatalf("mkdir src snapshot: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(srcSnapshot, "file.txt"), []byte("ok"), 0600); err != nil {
|
|
t.Fatalf("write snapshot file: %v", err)
|
|
}
|
|
sha, err := worker.DirOverallSHA256Hex(srcSnapshot)
|
|
if err != nil {
|
|
t.Fatalf("DirOverallSHA256Hex: %v", err)
|
|
}
|
|
cacheDir := filepath.Join(dataDir, "snapshots", "sha256", sha)
|
|
if err := os.MkdirAll(filepath.Dir(cacheDir), 0750); err != nil {
|
|
t.Fatalf("mkdir cache parent: %v", err)
|
|
}
|
|
if err := os.Rename(srcSnapshot, cacheDir); err != nil {
|
|
t.Fatalf("rename into cache: %v", err)
|
|
}
|
|
|
|
tq, err := queue.NewTaskQueue(queue.Config{RedisAddr: s.Addr(), MetricsFlushInterval: 5 * time.Millisecond})
|
|
if err != nil {
|
|
t.Fatalf("NewTaskQueue: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = tq.Close() })
|
|
|
|
task := &queue.Task{
|
|
ID: "task-1",
|
|
JobName: "job-1",
|
|
Status: "queued",
|
|
Priority: 10,
|
|
CreatedAt: time.Now().UTC(),
|
|
SnapshotID: "snap-1",
|
|
Metadata: map[string]string{
|
|
"snapshot_sha256": "sha256:" + sha,
|
|
},
|
|
}
|
|
if err := tq.AddTask(task); err != nil {
|
|
t.Fatalf("AddTask: %v", err)
|
|
}
|
|
|
|
cfg := &worker.Config{
|
|
WorkerID: "worker-1",
|
|
BasePath: base,
|
|
DataDir: dataDir,
|
|
PrewarmEnabled: true,
|
|
AutoFetchData: false,
|
|
LocalMode: true,
|
|
PollInterval: 1,
|
|
MaxWorkers: 1,
|
|
DatasetCacheTTL: 30 * time.Minute,
|
|
}
|
|
w := workertest.NewTestWorkerWithQueue(cfg, tq)
|
|
|
|
ok, err := w.PrewarmNextOnce(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("PrewarmNextOnce: %v", err)
|
|
}
|
|
if !ok {
|
|
t.Fatalf("expected ok=true")
|
|
}
|
|
|
|
// Empty the queue and run again. This should not delete the prewarm state; it should
|
|
// simply cancel its internal state and let the Redis TTL expire naturally.
|
|
_, _ = tq.GetNextTask() // drain the only queued task
|
|
_, _ = w.PrewarmNextOnce(context.Background())
|
|
|
|
state, err := tq.GetWorkerPrewarmState(cfg.WorkerID)
|
|
if err != nil {
|
|
t.Fatalf("GetWorkerPrewarmState: %v", err)
|
|
}
|
|
if state == nil {
|
|
t.Fatalf("expected prewarm state to remain present when queue is empty")
|
|
}
|
|
}
|
|
|
|
func TestStageSnapshotFromPath_UsesPrewarm(t *testing.T) {
|
|
base := t.TempDir()
|
|
taskID := "task-1"
|
|
jobDir := filepath.Join(base, "pending", "job-1", "run")
|
|
if err := os.MkdirAll(jobDir, 0750); err != nil {
|
|
t.Fatalf("mkdir jobDir: %v", err)
|
|
}
|
|
|
|
// Create a prewarmed snapshot directory.
|
|
prewarmSrc := filepath.Join(base, ".prewarm", "snapshots", taskID)
|
|
if err := os.MkdirAll(prewarmSrc, 0750); err != nil {
|
|
t.Fatalf("mkdir prewarm parent: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(prewarmSrc, "file.txt"), []byte("prewarmed"), 0600); err != nil {
|
|
t.Fatalf("write prewarm file: %v", err)
|
|
}
|
|
|
|
// Call stageSnapshotFromPath with a dummy srcPath; it should prefer the prewarmed dir.
|
|
dummySrc := filepath.Join(base, "unused")
|
|
if err := worker.StageSnapshotFromPath(base, taskID, dummySrc, jobDir); err != nil {
|
|
t.Fatalf("StageSnapshotFromPath: %v", err)
|
|
}
|
|
|
|
// Verify the prewarmed content was renamed into the job snapshot dir.
|
|
dstFile := filepath.Join(jobDir, "snapshot", "file.txt")
|
|
if _, err := os.Stat(dstFile); err != nil {
|
|
t.Fatalf("expected prewarmed file in job snapshot dir: %v", err)
|
|
}
|
|
if _, err := os.Stat(prewarmSrc); err == nil {
|
|
t.Fatalf("expected prewarm src to be moved (rename) not copied")
|
|
}
|
|
|
|
// Verify the content is correct.
|
|
got, err := os.ReadFile(dstFile)
|
|
if err != nil {
|
|
t.Fatalf("read dst file: %v", err)
|
|
}
|
|
if string(got) != "prewarmed" {
|
|
t.Fatalf("expected content 'prewarmed', got %q", string(got))
|
|
}
|
|
}
|
|
|
|
func TestStageSnapshotFromPath_FallsBackToCopy_WhenNoPrewarm(t *testing.T) {
|
|
base := t.TempDir()
|
|
taskID := "task-1"
|
|
jobDir := filepath.Join(base, "pending", "job-1", "run")
|
|
if err := os.MkdirAll(jobDir, 0750); err != nil {
|
|
t.Fatalf("mkdir jobDir: %v", err)
|
|
}
|
|
|
|
// Create a source snapshot dir (simulating ResolveSnapshot result).
|
|
src := filepath.Join(base, "src")
|
|
if err := os.MkdirAll(src, 0750); err != nil {
|
|
t.Fatalf("mkdir src: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(src, "file.txt"), []byte("source"), 0600); err != nil {
|
|
t.Fatalf("write src file: %v", err)
|
|
}
|
|
|
|
// No prewarm dir exists; should copy from src.
|
|
if err := worker.StageSnapshotFromPath(base, taskID, src, jobDir); err != nil {
|
|
t.Fatalf("StageSnapshotFromPath: %v", err)
|
|
}
|
|
|
|
dstFile := filepath.Join(jobDir, "snapshot", "file.txt")
|
|
if _, err := os.Stat(dstFile); err != nil {
|
|
t.Fatalf("expected file in job snapshot dir: %v", err)
|
|
}
|
|
got, err := os.ReadFile(dstFile)
|
|
if err != nil {
|
|
t.Fatalf("read dst file: %v", err)
|
|
}
|
|
if string(got) != "source" {
|
|
t.Fatalf("expected content 'source', got %q", string(got))
|
|
}
|
|
// Verify src is still present (copy, not move).
|
|
if _, err := os.Stat(filepath.Join(src, "file.txt")); err != nil {
|
|
t.Fatalf("expected src file to remain after copy")
|
|
}
|
|
}
|