orca/internal/transport/dispatch_test.go

package transport

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

type mockDispatcher struct {
	jobID     string
	state     string
	submitErr error
	statusErr error
	submits   int
	statuses  int
	lastSpec  []byte
}

func (m *mockDispatcher) LocalSubmit(ctx context.Context, spec []byte) (string, error) {
	m.submits++
	m.lastSpec = spec
	if m.submitErr != nil {
		return "", m.submitErr
	}
	if m.jobID == "" {
		return "job-123", nil
	}
	return m.jobID, nil
}

func (m *mockDispatcher) LocalStatus(ctx context.Context, jobID string) (string, error) {
	m.statuses++
	if m.statusErr != nil {
		return "", m.statusErr
	}
	if m.state == "" {
		return "running", nil
	}
	return m.state, nil
}

func TestSubmitHandler_Success(t *testing.T) {
	d := &mockDispatcher{}
	h := NewSubmitHandler(d, nil)
	body := bytes.NewReader([]byte(`{"spec":"{}"}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Submit", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("status = %d, want 200", w.Code)
	}
	var resp SubmitResponse
	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if resp.JobID != "job-123" {
		t.Errorf("JobID = %q, want job-123", resp.JobID)
	}
	if d.submits != 1 {
		t.Errorf("submits = %d, want 1", d.submits)
	}
}

func TestSubmitHandler_IdempotencyReplay(t *testing.T) {
	d := &mockDispatcher{}
	store := NewIdempotencyStore()
	store.Put("key-1", "job-existing")
	h := NewSubmitHandler(d, store)
	body := bytes.NewReader([]byte(`{"spec":"{}"}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Submit", body)
	req.Header.Set(IdempotencyHeader, "key-1")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("status = %d, want 200", w.Code)
	}
	var resp SubmitResponse
	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if resp.JobID != "job-existing" {
		t.Errorf("JobID = %q, want job-existing (replay)", resp.JobID)
	}
	if d.submits != 0 {
		t.Errorf("submits = %d, want 0 (replayed from store)", d.submits)
	}
}

func TestSubmitHandler_IdempotencyStores(t *testing.T) {
	d := &mockDispatcher{}
	store := NewIdempotencyStore()
	h := NewSubmitHandler(d, store)
	body := bytes.NewReader([]byte(`{"spec":"{}"}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Submit", body)
	req.Header.Set(IdempotencyHeader, "key-2")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200", w.Code)
	}
	if got, ok := store.Get("key-2"); !ok || got != "job-123" {
		t.Errorf("store.Get(key-2) = (%q, %v), want (job-123, true)", got, ok)
	}
}

func TestSubmitHandler_BadMethod(t *testing.T) {
	h := NewSubmitHandler(&mockDispatcher{}, nil)
	req := httptest.NewRequest(http.MethodGet, "/orca.v1.Dispatch/Submit", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusMethodNotAllowed {
		t.Errorf("status = %d, want 405", w.Code)
	}
}

func TestSubmitHandler_BadBody(t *testing.T) {
	h := NewSubmitHandler(&mockDispatcher{}, nil)
	body := strings.NewReader("{not json")
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Submit", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusBadRequest {
		t.Errorf("status = %d, want 400", w.Code)
	}
}

func TestSubmitHandler_EmptySpec(t *testing.T) {
	h := NewSubmitHandler(&mockDispatcher{}, nil)
	body := bytes.NewReader([]byte(`{}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Submit", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusBadRequest {
		t.Errorf("status = %d, want 400", w.Code)
	}
}

func TestSubmitHandler_DispatcherError(t *testing.T) {
	d := &mockDispatcher{submitErr: errors.New("boom")}
	h := NewSubmitHandler(d, nil)
	body := bytes.NewReader([]byte(`{"spec":"{}"}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Submit", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusInternalServerError {
		t.Errorf("status = %d, want 500", w.Code)
	}
}

func TestStatusHandler_Success(t *testing.T) {
	d := &mockDispatcher{state: "complete"}
	h := NewStatusHandler(d)
	body := bytes.NewReader([]byte(`{"job_id":"job-1"}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Status", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Errorf("status = %d, want 200", w.Code)
	}
	var resp StatusResponse
	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if resp.State != "complete" {
		t.Errorf("State = %q, want complete", resp.State)
	}
}

func TestStatusHandler_BadMethod(t *testing.T) {
	h := NewStatusHandler(&mockDispatcher{})
	req := httptest.NewRequest(http.MethodGet, "/orca.v1.Dispatch/Status", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusMethodNotAllowed {
		t.Errorf("status = %d, want 405", w.Code)
	}
}

func TestStatusHandler_BadBody(t *testing.T) {
	h := NewStatusHandler(&mockDispatcher{})
	body := strings.NewReader("nope")
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Status", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusBadRequest {
		t.Errorf("status = %d, want 400", w.Code)
	}
}

func TestStatusHandler_EmptyJobID(t *testing.T) {
	h := NewStatusHandler(&mockDispatcher{})
	body := bytes.NewReader([]byte(`{"job_id":""}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Status", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusBadRequest {
		t.Errorf("status = %d, want 400", w.Code)
	}
}

func TestStatusHandler_DispatcherError(t *testing.T) {
	d := &mockDispatcher{statusErr: errors.New("not found")}
	h := NewStatusHandler(d)
	body := bytes.NewReader([]byte(`{"job_id":"job-x"}`))
	req := httptest.NewRequest(http.MethodPost, "/orca.v1.Dispatch/Status", body)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusNotFound {
		t.Errorf("status = %d, want 404", w.Code)
	}
}

func TestNewDispatchClient(t *testing.T) {
	dir := t.TempDir()
	ca, err := security.CAInit(dir, "orca-test-ca")
	if err != nil {
		t.Fatalf("CAInit: %v", err)
	}
	caPath := filepath.Join(dir, security.CACertFile)
	_ = ca
	c, err := NewDispatchClient(caPath, "localhost", "https://localhost:8443")
	if err != nil {
		t.Fatalf("NewDispatchClient: %v", err)
	}
	if c == nil {
		t.Fatal("client is nil")
	}
	if c.PeerAddr != "https://localhost:8443" {
		t.Errorf("PeerAddr = %q, want https://localhost:8443", c.PeerAddr)
	}
}

func TestNewDispatchClient_EmptyCAPath(t *testing.T) {
	_, err := NewDispatchClient("", "localhost", "https://localhost:8443")
	if err == nil {
		t.Fatal("expected error for empty caPath")
	}
}

func TestNewDispatchClient_EmptyServerName(t *testing.T) {
	dir := t.TempDir()
	_, err := security.CAInit(dir, "orca-test-ca")
	caPath := filepath.Join(dir, security.CACertFile)
	_, err = NewDispatchClient(caPath, "", "https://localhost:8443")
	if err == nil {
		t.Fatal("expected error for empty serverName")
	}
}

func TestDispatchClient_InvalidURL(t *testing.T) {
	dir := t.TempDir()
	_, err := security.CAInit(dir, "orca-test-ca")
	if err != nil {
		t.Fatalf("CAInit: %v", err)
	}
	caPath := filepath.Join(dir, security.CACertFile)
	c, err := NewDispatchClient(caPath, "localhost", "http://127.0.0.1:1")
	if err != nil {
		t.Fatalf("NewDispatchClient: %v", err)
	}
	_, err = c.Status(context.Background(), "job-1")
	if err == nil {
		t.Fatal("expected error for connection refused")
	}
}

func TestDispatchClient_Status_HTTPError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeError(w, http.StatusInternalServerError, "boom")
	}))
	defer srv.Close()

	dc := &DispatchClient{
		HTTP:     &MTLSClient{http: &http.Client{Timeout: 5 * time.Second}},
		PeerAddr: srv.URL,
	}
	_, err := dc.Status(context.Background(), "job-1")
	if err == nil {
		t.Fatal("expected error for 500 status")
	}
}

func TestDispatchClient_Status_DecodeError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("{not valid json"))
	}))
	defer srv.Close()

	dc := &DispatchClient{
		HTTP:     &MTLSClient{http: &http.Client{Timeout: 5 * time.Second}},
		PeerAddr: srv.URL,
	}
	_, err := dc.Status(context.Background(), "job-1")
	if err == nil {
		t.Fatal("expected decode error")
	}
}

func TestDispatchClient_Status_Success(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeError(w, http.StatusMethodNotAllowed, "method")
			return
		}
		body, _ := io.ReadAll(r.Body)
		var req StatusRequest
		_ = json.Unmarshal(body, &req)
		if req.JobID != "job-9" {
			writeError(w, http.StatusBadRequest, "bad job_id")
			return
		}
		writeJSON(w, http.StatusOK, StatusResponse{JobID: "job-9", NodeID: "self", State: "complete"})
	}))
	defer srv.Close()

	dc := &DispatchClient{
		HTTP:     &MTLSClient{http: &http.Client{Timeout: 5 * time.Second}},
		PeerAddr: srv.URL,
	}
	resp, err := dc.Status(context.Background(), "job-9")
	if err != nil {
		t.Fatalf("Status: %v", err)
	}
	if resp.JobID != "job-9" {
		t.Errorf("JobID = %q, want job-9", resp.JobID)
	}
	if resp.State != "complete" {
		t.Errorf("State = %q, want complete", resp.State)
	}
}

func TestDispatchClient_Submit_Success(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeError(w, http.StatusMethodNotAllowed, "method")
			return
		}
		writeJSON(w, http.StatusOK, SubmitResponse{JobID: "job-submit-1", NodeID: "peer-1"})
	}))
	defer srv.Close()

	dc := &DispatchClient{
		HTTP:     &MTLSClient{http: &http.Client{Timeout: 5 * time.Second}},
		PeerAddr: srv.URL,
	}
	resp, err := dc.Submit(context.Background(), []byte("spec"), "idem-key-1")
	if err != nil {
		t.Fatalf("Submit: %v", err)
	}
	if resp.JobID != "job-submit-1" {
		t.Errorf("JobID = %q, want job-submit-1", resp.JobID)
	}
}

func TestDispatchClient_Submit_NonIdempotentTransientBails(t *testing.T) {
	calls := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		writeError(w, http.StatusServiceUnavailable, "unavailable")
	}))
	defer srv.Close()

	dc := &DispatchClient{
		HTTP:     &MTLSClient{http: &http.Client{Timeout: 5 * time.Second}},
		PeerAddr: srv.URL,
	}
	_, err := dc.Submit(context.Background(), []byte("spec"), "")
	if err == nil {
		t.Fatal("expected error")
	}
	if calls != 1 {
		t.Errorf("calls = %d, want 1 (no key, no retry)", calls)
	}
}

func TestBytesReader(t *testing.T) {
	r := bytesReader([]byte("hello"))
	buf := make([]byte, 5)
	n, err := r.Read(buf)
	if n != 5 || err != nil || string(buf) != "hello" {
		t.Errorf("Read: n=%d err=%v buf=%q", n, err, buf)
	}
	n, err = r.Read(buf)
	if n != 0 || err == nil {
		t.Errorf("Read past end: n=%d err=%v, want error", n, err)
	}
	if err := r.Close(); err != nil {
		t.Errorf("Close: %v", err)
	}
}