orca/internal/daemon/dispatch_test.go

// Package daemon — dispatch_test.go exercises the orca.v1.Dispatch
// round-trip end-to-end: a SubmitHandler is mounted on a test server
// and a DispatchClient dials it. The test asserts the spec flows
// through, the job ID is returned, and dedupe (X-Orca-Idempotency-Key)
// works.
package daemon

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

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

// stubDispatcher is a transport.Dispatcher for tests. It records
// every Submit and Status call and returns deterministic responses.
type stubDispatcher struct {
	mu         sync.Mutex
	submits    [][]byte
	statuses   []string
	nextJobID  int
	failSubmit bool
}

func (s *stubDispatcher) LocalSubmit(_ context.Context, spec []byte) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.failSubmit {
		return "", fmt.Errorf("submit failed (test)")
	}
	cp := make([]byte, len(spec))
	copy(cp, spec)
	s.submits = append(s.submits, cp)
	s.nextJobID++
	return fmt.Sprintf("job-%d", s.nextJobID), nil
}

func (s *stubDispatcher) LocalStatus(_ context.Context, jobID string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.statuses = append(s.statuses, jobID)
	return "running", nil
}

func TestDispatchRoundTrip(t *testing.T) {
	stub := &stubDispatcher{}
	dedupe := transport.NewIdempotencyStore()
	handlers := NewDispatchHandlers(stub, dedupe)

	mux := http.NewServeMux()
	handlers.Mount(mux)
	ts := httptest.NewServer(mux)
	t.Cleanup(ts.Close)

	// Submit a spec wrapped in the SubmitRequest envelope.
	// The wire format is {"spec": <json.RawMessage>}; the inner
	// spec is opaque to the dispatch service and is parsed by the
	// local executor downstream.
	inner := []byte(`{"name":"hello","command":"/bin/echo","args":["hi"],"env":[]}`)
	wire, _ := json.Marshal(transport.SubmitRequest{Spec: inner})
	resp, err := http.Post(ts.URL+"/orca.v1.Dispatch/Submit", "application/json", bytes.NewReader(wire))
	if err != nil {
		t.Fatalf("Submit: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Submit status: got %d, want 200", resp.StatusCode)
	}
	body, _ := io.ReadAll(resp.Body)
	var sr transport.SubmitResponse
	if err := json.Unmarshal(body, &sr); err != nil {
		t.Fatalf("decode Submit response: %v", err)
	}
	if sr.JobID == "" {
		t.Fatal("Submit response missing job_id")
	}
	if len(stub.submits) != 1 {
		t.Errorf("LocalSubmit calls: got %d, want 1", len(stub.submits))
	}

	// Status query.
	statusReq := transport.StatusRequest{JobID: sr.JobID}
	body2, _ := json.Marshal(statusReq)
	resp2, err := http.Post(ts.URL+"/orca.v1.Dispatch/Status", "application/json", bytes.NewReader(body2))
	if err != nil {
		t.Fatalf("Status: %v", err)
	}
	defer resp2.Body.Close()
	if resp2.StatusCode != http.StatusOK {
		t.Fatalf("Status code: got %d, want 200", resp2.StatusCode)
	}
	var stResp transport.StatusResponse
	if err := json.NewDecoder(resp2.Body).Decode(&stResp); err != nil {
		t.Fatalf("decode Status: %v", err)
	}
	if stResp.State != "running" {
		t.Errorf("Status.State: got %q, want running", stResp.State)
	}
}

func TestDispatchIdempotencyDedupe(t *testing.T) {
	stub := &stubDispatcher{}
	dedupe := transport.NewIdempotencyStore()
	handlers := NewDispatchHandlers(stub, dedupe)

	mux := http.NewServeMux()
	handlers.Mount(mux)
	ts := httptest.NewServer(mux)
	t.Cleanup(ts.Close)

	inner := []byte(`{"name":"hello","command":"/bin/echo","args":["hi"]}`)
	wire, _ := json.Marshal(transport.SubmitRequest{Spec: inner})
	post := func() string {
		req, _ := http.NewRequest(http.MethodPost, ts.URL+"/orca.v1.Dispatch/Submit", bytes.NewReader(wire))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set(transport.IdempotencyHeader, "key-42")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatalf("Submit: %v", err)
		}
		defer resp.Body.Close()
		b, _ := io.ReadAll(resp.Body)
		return string(b)
	}

	// First call: real submit, LocalSubmit invoked.
	first := post()
	var sr1 transport.SubmitResponse
	if err := json.Unmarshal([]byte(first), &sr1); err != nil {
		t.Fatalf("decode 1: %v", err)
	}
	if len(stub.submits) != 1 {
		t.Errorf("after first call: submits=%d, want 1", len(stub.submits))
	}

	// Second call: same key, dedupe replay.
	second := post()
	var sr2 transport.SubmitResponse
	if err := json.Unmarshal([]byte(second), &sr2); err != nil {
		t.Fatalf("decode 2: %v", err)
	}
	if sr1.JobID != sr2.JobID {
		t.Errorf("dedupe: first=%s, second=%s (should match)", sr1.JobID, sr2.JobID)
	}
	if len(stub.submits) != 1 {
		t.Errorf("after second call: submits=%d, want 1 (dedupe)", len(stub.submits))
	}
}

func TestDispatchSubmitValidation(t *testing.T) {
	stub := &stubDispatcher{}
	handlers := NewDispatchHandlers(stub, transport.NewIdempotencyStore())
	mux := http.NewServeMux()
	handlers.Mount(mux)
	ts := httptest.NewServer(mux)
	t.Cleanup(ts.Close)

	// Empty spec: 400.
	resp, _ := http.Post(ts.URL+"/orca.v1.Dispatch/Submit", "application/json", bytes.NewReader([]byte(`{}`)))
	if resp.StatusCode != http.StatusBadRequest {
		t.Errorf("empty spec: status=%d, want 400", resp.StatusCode)
	}
	resp.Body.Close()

	// GET instead of POST: 405.
	resp2, _ := http.Get(ts.URL + "/orca.v1.Dispatch/Submit")
	if resp2.StatusCode != http.StatusMethodNotAllowed {
		t.Errorf("GET: status=%d, want 405", resp2.StatusCode)
	}
	resp2.Body.Close()
}