orca/internal/engine/dispatcher_test.go

package engine

import (
	"context"
	"errors"
	"path/filepath"
	"testing"

	"git.cloudinit.dev/coreci/orca/internal/store"
)

type mockExecutor struct {
	submitFn  func(ctx context.Context, spec []byte) (string, error)
	statusFn  func(ctx context.Context, jobID string) (string, error)
	submitted bool
}

func (m *mockExecutor) Submit(ctx context.Context, spec []byte) (string, error) {
	m.submitted = true
	if m.submitFn != nil {
		return m.submitFn(ctx, spec)
	}
	return "mock-job-id", nil
}

func (m *mockExecutor) Status(ctx context.Context, jobID string) (string, error) {
	if m.statusFn != nil {
		return m.statusFn(ctx, jobID)
	}
	return "complete", nil
}

func newTestDispatcher(t *testing.T, exec LocalExecutor) (*Dispatcher, *store.CapacityRepo, func()) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "test.db")
	db, err := store.Open(path)
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	capRepo := store.NewCapacityRepo(db)
	peers := NewPeerRegistry()
	d := NewDispatcher(nil, capRepo, peers, exec)
	return d, capRepo, func() { _ = db.Close() }
}

func TestDispatcher_Submit_EmptySpec(t *testing.T) {
	d, _, cleanup := newTestDispatcher(t, &mockExecutor{})
	defer cleanup()
	_, _, err := d.Submit(context.Background(), "", nil, "")
	if err == nil {
		t.Fatal("Submit: expected error for empty spec, got nil")
	}
}

func TestDispatcher_Submit_IdempotencyHit(t *testing.T) {
	exec := &mockExecutor{}
	d, _, cleanup := newTestDispatcher(t, exec)
	defer cleanup()

	d.Dedupe().Put("key-1", "cached-job-id")
	spec := []byte(`{"cpu_millicores":100,"memory_mib":64,"disk_mib":64}`)
	jobID, nodeID, err := d.Submit(context.Background(), "", spec, "key-1")
	if err != nil {
		t.Fatalf("Submit: %v", err)
	}
	if jobID != "cached-job-id" {
		t.Errorf("jobID: got %q, want cached-job-id", jobID)
	}
	if nodeID != "self" {
		t.Errorf("nodeID: got %q, want self", nodeID)
	}
	if exec.submitted {
		t.Error("executor was called on idempotency hit; should have been short-circuited")
	}
}

func TestDispatcher_Submit_LocalCapacity(t *testing.T) {
	exec := &mockExecutor{
		submitFn: func(ctx context.Context, spec []byte) (string, error) {
			return "local-job-id", nil
		},
	}
	d, capRepo, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	ctx := context.Background()
	if err := capRepo.Upsert(ctx, &store.NodeCapacity{
		NodeID:        "self",
		CPUMillicores: 4000,
		MemoryMiB:     4096,
		DiskMiB:       4096,
	}); err != nil {
		t.Fatalf("Upsert capacity: %v", err)
	}
	spec := []byte(`{"cpu_millicores":100,"memory_mib":64,"disk_mib":64}`)
	jobID, nodeID, err := d.Submit(ctx, "", spec, "")
	if err != nil {
		t.Fatalf("Submit: %v", err)
	}
	if jobID != "local-job-id" {
		t.Errorf("jobID: got %q, want local-job-id", jobID)
	}
	if nodeID != "self" {
		t.Errorf("nodeID: got %q, want self", nodeID)
	}
	if !exec.submitted {
		t.Error("executor was not called for local-capacity path")
	}
}

func TestDispatcher_Submit_ExplicitTarget(t *testing.T) {
	exec := &mockExecutor{}
	d, _, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	spec := []byte(`{"cpu_millicores":100,"memory_mib":64,"disk_mib":64}`)
	_, _, err := d.Submit(context.Background(), "nodeA", spec, "")
	if err == nil {
		t.Fatal("Submit with explicit target nodeA (no peer): expected error, got nil")
	}
}

func TestDispatcher_Submit_NoPeers(t *testing.T) {
	exec := &mockExecutor{}
	d, capRepo, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	ctx := context.Background()
	if err := capRepo.Upsert(ctx, &store.NodeCapacity{
		NodeID:        "self",
		CPUMillicores: 0,
		MemoryMiB:     0,
		DiskMiB:       0,
	}); err != nil {
		t.Fatalf("Upsert: %v", err)
	}
	spec := []byte(`{"cpu_millicores":1000,"memory_mib":1024,"disk_mib":1024}`)
	_, _, err := d.Submit(ctx, "", spec, "")
	if err == nil {
		t.Fatal("Submit: expected error when no peers and no local capacity, got nil")
	}
}

func TestDispatcher_LocalSubmit(t *testing.T) {
	exec := &mockExecutor{
		submitFn: func(ctx context.Context, spec []byte) (string, error) {
			return "ls-job", nil
		},
	}
	d, _, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	jobID, err := d.LocalSubmit(context.Background(), []byte(`{"command":"/bin/true"}`))
	if err != nil {
		t.Fatalf("LocalSubmit: %v", err)
	}
	if jobID != "ls-job" {
		t.Errorf("LocalSubmit: got %q, want ls-job", jobID)
	}
	if !exec.submitted {
		t.Error("LocalSubmit: executor.Submit not called")
	}
}

func TestDispatcher_LocalStatus(t *testing.T) {
	exec := &mockExecutor{
		statusFn: func(ctx context.Context, jobID string) (string, error) {
			if jobID == "known" {
				return "running", nil
			}
			return "", errors.New("not found")
		},
	}
	d, _, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	st, err := d.LocalStatus(context.Background(), "known")
	if err != nil {
		t.Fatalf("LocalStatus: %v", err)
	}
	if st != "running" {
		t.Errorf("LocalStatus: got %q, want running", st)
	}
	if _, err := d.LocalStatus(context.Background(), "missing"); err == nil {
		t.Error("LocalStatus: expected error for missing job, got nil")
	}
}

func TestDispatcher_LocalSubmit_NilExecutor(t *testing.T) {
	d := NewDispatcher(nil, nil, NewPeerRegistry(), nil)
	if _, err := d.LocalSubmit(context.Background(), []byte(`{}`)); err == nil {
		t.Error("LocalSubmit with nil executor: expected error, got nil")
	}
	if _, err := d.LocalStatus(context.Background(), "x"); err == nil {
		t.Error("LocalStatus with nil executor: expected error, got nil")
	}
}

func TestParseInlineSpec(t *testing.T) {
	spec, err := parseInlineSpec([]byte(`{"cpu_millicores":500,"memory_mib":256,"disk_mib":128}`))
	if err != nil {
		t.Fatalf("parseInlineSpec: %v", err)
	}
	if spec.CPUMillicores != 500 || spec.MemoryMiB != 256 || spec.DiskMiB != 128 {
		t.Errorf("parseInlineSpec: got %+v, want cpu=500 mem=256 disk=128", spec)
	}
	if _, err := parseInlineSpec([]byte(`{bad json`)); err == nil {
		t.Fatal("parseInlineSpec: expected error for malformed JSON, got nil")
	}
}

func TestDispatcher_Submit_BadSpec(t *testing.T) {
	d, _, cleanup := newTestDispatcher(t, &mockExecutor{})
	defer cleanup()
	_, _, err := d.Submit(context.Background(), "", []byte(`{bad json`), "")
	if err == nil {
		t.Fatal("expected error for malformed spec")
	}
}

func TestDispatcher_Submit_ExplicitTargetNoPeerRegistry(t *testing.T) {
	d := NewDispatcher(nil, nil, nil, &mockExecutor{})
	_, _, err := d.Submit(context.Background(), "nodeX", []byte(`{"cpu_millicores":100,"memory_mib":64,"disk_mib":64}`), "")
	if err == nil {
		t.Fatal("expected error for explicit target with no peer registry")
	}
}

func TestDispatcher_Submit_ExplicitTargetPeerNotFound(t *testing.T) {
	d, _, cleanup := newTestDispatcher(t, &mockExecutor{})
	defer cleanup()
	_, _, err := d.Submit(context.Background(), "ghost", []byte(`{"cpu_millicores":100,"memory_mib":64,"disk_mib":64}`), "")
	if err == nil {
		t.Fatal("expected error for target not in registry")
	}
}

func TestDispatcher_Submit_PickPeerMissingCA(t *testing.T) {
	exec := &mockExecutor{}
	d, capRepo, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	ctx := context.Background()
	if err := capRepo.Upsert(ctx, &store.NodeCapacity{
		NodeID:        "self",
		CPUMillicores: 0,
		MemoryMiB:     0,
		DiskMiB:       0,
	}); err != nil {
		t.Fatalf("Upsert: %v", err)
	}
	if err := d.peers.Add(&Peer{
		NodeID:   "peer-1",
		Address:  "127.0.0.1:1",
		Capacity: &store.NodeCapacity{NodeID: "peer-1", CPUMillicores: 4000, MemoryMiB: 4096, DiskMiB: 4096},
	}); err != nil {
		t.Fatalf("Add peer: %v", err)
	}
	spec := []byte(`{"cpu_millicores":100,"memory_mib":64,"disk_mib":64}`)
	_, _, err := d.Submit(ctx, "", spec, "idem-peer-1")
	if err == nil {
		t.Fatal("expected error (peer missing CA/servername)")
	}
}

func TestDispatcher_Submit_NoPeerRegistry(t *testing.T) {
	d := NewDispatcher(nil, nil, nil, &mockExecutor{})
	spec := []byte(`{"cpu_millicores":1000,"memory_mib":1024,"disk_mib":1024}`)
	_, _, err := d.Submit(context.Background(), "", spec, "")
	if err == nil {
		t.Fatal("expected error for no peer registry and no capacity repo")
	}
}

func TestDispatcher_Submit_NilCapacityFallsThrough(t *testing.T) {
	exec := &mockExecutor{}
	d := NewDispatcher(nil, nil, NewPeerRegistry(), exec)
	spec := []byte(`{"cpu_millicores":1000,"memory_mib":1024,"disk_mib":1024}`)
	_, _, err := d.Submit(context.Background(), "", spec, "")
	if err == nil {
		t.Fatal("expected error when capacity repo is nil and no peers")
	}
}

func TestDispatcher_Submit_AllPeersFailsPickNode(t *testing.T) {
	exec := &mockExecutor{}
	d, capRepo, cleanup := newTestDispatcher(t, exec)
	defer cleanup()
	ctx := context.Background()
	if err := capRepo.Upsert(ctx, &store.NodeCapacity{
		NodeID:        "self",
		CPUMillicores: 0,
		MemoryMiB:     0,
		DiskMiB:       0,
	}); err != nil {
		t.Fatalf("Upsert: %v", err)
	}
	if err := d.peers.Add(&Peer{
		NodeID:   "peer-tiny",
		Address:  "127.0.0.1:1",
		Capacity: &store.NodeCapacity{NodeID: "peer-tiny", CPUMillicores: 10, MemoryMiB: 10, DiskMiB: 10},
	}); err != nil {
		t.Fatalf("Add peer: %v", err)
	}
	spec := []byte(`{"cpu_millicores":1000,"memory_mib":1024,"disk_mib":1024}`)
	_, _, err := d.Submit(ctx, "", spec, "")
	if err == nil {
		t.Fatal("expected error when no peer can fit")
	}
}