fetch_ml/tests/unit/worker/test_helpers.go

// Package worker_test provides test helpers for the worker package
package worker_test

import (
	"context"
	"log/slog"

	"github.com/jfraeys/fetch_ml/internal/jupyter"
	"github.com/jfraeys/fetch_ml/internal/logging"
	"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/lifecycle"
)

// 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{}

	return &worker.Worker{
		ID:      "test-worker",
		Config:  cfg,
		Logger:  logger,
		Metrics: metricsObj,
		Health:  lifecycle.NewHealthMonitor(),
	}
}

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

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

// ResolveDatasets resolves dataset paths for a task.
func ResolveDatasets(ctx context.Context, w *worker.Worker, task *queue.Task) ([]string, error) {
	if task.DatasetSpecs == nil {
		return nil, nil
	}

	dataDir := w.Config.DataDir
	if dataDir == "" {
		dataDir = "/tmp/data"
	}

	var paths []string
	for _, spec := range task.DatasetSpecs {
		path := dataDir + "/" + spec.Name
		paths = append(paths, path)
	}

	return paths, nil
}