fetch_ml/internal/workertest/worker.go

// Package workertest provides test helpers for the worker package.
// This package is only intended for use in tests and is separate from
// production code to maintain clean separation of concerns.
package workertest

import (
	"log/slog"
	"strings"
	"time"

	"github.com/jfraeys/fetch_ml/internal/logging"
	"github.com/jfraeys/fetch_ml/internal/manifest"
	"github.com/jfraeys/fetch_ml/internal/metrics"
	"github.com/jfraeys/fetch_ml/internal/queue"
	"github.com/jfraeys/fetch_ml/internal/worker"
	"github.com/jfraeys/fetch_ml/internal/worker/executor"
	"github.com/jfraeys/fetch_ml/internal/worker/lifecycle"
)

// SimpleManifestWriter is a basic ManifestWriter implementation for testing
type SimpleManifestWriter struct{}

func (w *SimpleManifestWriter) Upsert(dir string, task *queue.Task, mutate func(*manifest.RunManifest)) {
	// Try to load existing manifest, or create new one
	m, err := manifest.LoadFromDir(dir)
	if err != nil {
		m = w.BuildInitial(task, "")
	}
	mutate(m)
	_ = m.WriteToDir(dir)
}

func (w *SimpleManifestWriter) BuildInitial(task *queue.Task, podmanImage string) *manifest.RunManifest {
	m := manifest.NewRunManifest(
		"run-"+task.ID,
		task.ID,
		task.JobName,
		time.Now().UTC(),
	)
	m.CommitID = task.Metadata["commit_id"]
	m.DepsManifestName = task.Metadata["deps_manifest_name"]
	return m
}

// NewTestWorker creates a minimal Worker for testing purposes.
// It initializes only the fields needed for unit tests.
func NewTestWorker(cfg *worker.Config) *worker.Worker {
	if cfg == nil {
		cfg = &worker.Config{}
	}

	logger := logging.NewLogger(slog.LevelInfo, false)
	metricsObj := &metrics.Metrics{}

	// Create executors and runner for testing
	writer := &SimpleManifestWriter{}
	localExecutor := executor.NewLocalExecutor(logger, writer)
	containerExecutor := executor.NewContainerExecutor(
		logger,
		nil,
		executor.ContainerConfig{
			PodmanImage: cfg.PodmanImage,
			BasePath:    cfg.BasePath,
		},
	)
	jobRunner := executor.NewJobRunner(
		localExecutor,
		containerExecutor,
		writer,
		logger,
	)

	return &worker.Worker{
		ID:      cfg.WorkerID,
		Config:  cfg,
		Logger:  logger,
		Metrics: metricsObj,
		Health:  lifecycle.NewHealthMonitor(),
		Runner:  jobRunner,
	}
}

// NewTestWorkerWithQueue creates a test Worker with a queue client.
func NewTestWorkerWithQueue(cfg *worker.Config, queueClient queue.Backend) *worker.Worker {
	w := NewTestWorker(cfg)
	w.QueueClient = queueClient
	return w
}

// NewTestWorkerWithJupyter creates a test Worker with Jupyter manager.
func NewTestWorkerWithJupyter(cfg *worker.Config, jupyterMgr worker.JupyterManager) *worker.Worker {
	w := NewTestWorker(cfg)
	w.Jupyter = jupyterMgr
	return w
}

// NewTestWorkerWithRunner creates a test Worker with JobRunner initialized.
// Note: This creates a minimal runner for testing purposes.
func NewTestWorkerWithRunner(cfg *worker.Config) *worker.Worker {
	return NewTestWorker(cfg)
}

// NewTestWorkerWithRunLoop creates a test Worker with RunLoop initialized.
// Note: RunLoop requires proper queue client setup.
func NewTestWorkerWithRunLoop(cfg *worker.Config, queueClient queue.Backend) *worker.Worker {
	return NewTestWorkerWithQueue(cfg, queueClient)
}

// ResolveDatasets resolves dataset paths for a task.
// This version matches the test expectations for backwards compatibility.
// Priority: DatasetSpecs > Datasets > Args parsing
func ResolveDatasets(task *queue.Task) []string {
	if task == nil {
		return nil
	}

	// Priority 1: DatasetSpecs
	if len(task.DatasetSpecs) > 0 {
		var paths []string
		for _, spec := range task.DatasetSpecs {
			paths = append(paths, spec.Name)
		}
		return paths
	}

	// Priority 2: Datasets
	if len(task.Datasets) > 0 {
		return task.Datasets
	}

	// Priority 3: Parse from Args
	if task.Args != "" {
		// Simple parsing: --datasets a,b,c or --datasets a b c
		args := task.Args
		if idx := strings.Index(args, "--datasets"); idx != -1 {
			after := args[idx+len("--datasets "):]
			after = strings.TrimSpace(after)
			// Split by comma or space
			if strings.Contains(after, ",") {
				return strings.Split(after, ",")
			}
			parts := strings.Fields(after)
			if len(parts) > 0 {
				return parts
			}
		}
	}

	return nil
}