package httpserver import ( "bytes" "encoding/json" "net/http" "net/http/httptest" "testing" "github.com/toki/oto/services/core/internal/runnerregistry" otopb "github.com/toki/oto/services/core/oto" ) func TestHandleHealthz(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "/healthz", nil) rr := httptest.NewRecorder() handleHealthz(rr, req) if status := rr.Code; status != http.StatusOK { t.Errorf("handleHealthz returned wrong status code: got %v want %v", status, http.StatusOK) } expected := "OK" if body := rr.Body.String(); body != expected { t.Errorf("handleHealthz returned unexpected body: got %v want %v", body, expected) } } func TestHandleHealthz_MethodNotAllowed(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/healthz", nil) rr := httptest.NewRecorder() handleHealthz(rr, req) if status := rr.Code; status != http.StatusMethodNotAllowed { t.Errorf("handleHealthz for POST returned wrong status code: got %v want %v", status, http.StatusMethodNotAllowed) } } func TestHandleReadyz(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "/readyz", nil) rr := httptest.NewRecorder() handleReadyz(rr, req) if status := rr.Code; status != http.StatusOK { t.Errorf("handleReadyz returned wrong status code: got %v want %v", status, http.StatusOK) } expected := "OK" if body := rr.Body.String(); body != expected { t.Errorf("handleReadyz returned unexpected body: got %v want %v", body, expected) } } func TestHandleReadyz_MethodNotAllowed(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "/readyz", nil) rr := httptest.NewRecorder() handleReadyz(rr, req) if status := rr.Code; status != http.StatusMethodNotAllowed { t.Errorf("handleReadyz for POST returned wrong status code: got %v want %v", status, http.StatusMethodNotAllowed) } } func TestServerMux(t *testing.T) { registry := runnerregistry.New() server := NewServerWithRegistry(":0", registry) mux := server.httpServer.Handler.(*http.ServeMux) // Test healthz routing reqHealth := httptest.NewRequest(http.MethodGet, "/healthz", nil) rrHealth := httptest.NewRecorder() mux.ServeHTTP(rrHealth, reqHealth) if rrHealth.Code != http.StatusOK { t.Errorf("mux /healthz failed: got status %v", rrHealth.Code) } // Test readyz routing reqReady := httptest.NewRequest(http.MethodGet, "/readyz", nil) rrReady := httptest.NewRecorder() mux.ServeHTTP(rrReady, reqReady) if rrReady.Code != http.StatusOK { t.Errorf("mux /readyz failed: got status %v", rrReady.Code) } // Test fallback/not found routing reqNotFound := httptest.NewRequest(http.MethodGet, "/unknown", nil) rrNotFound := httptest.NewRecorder() mux.ServeHTTP(rrNotFound, reqNotFound) if rrNotFound.Code != http.StatusNotFound { t.Errorf("mux /unknown routing status code: got %v want %v", rrNotFound.Code, http.StatusNotFound) } } func TestHandleRunnerRegisterAcceptsAndStoresRunner(t *testing.T) { registry := runnerregistry.New() body := bytes.NewBufferString(`{ "enrollment_token":"token-123", "runner_id":"runner-123", "alias":"linux-build", "protocol_version":"oto.runner.v1", "capability":{"name":"oto-runner","version":"1.0.0"}, "command_catalog":{"command_types":["Shell","Git"]} }`) req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body) rr := httptest.NewRecorder() handleRunnerRegister(registry)(rr, req) if rr.Code != http.StatusOK { t.Fatalf("status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String()) } var response otopb.RegisterRunnerResponse if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil { t.Fatalf("decode response: %v", err) } if !response.GetAccepted() { t.Fatalf("accepted=false, reason=%q", response.GetRejectReason()) } record, ok := registry.Snapshot("runner-123") if !ok { t.Fatal("registered runner was not stored") } if record.Alias != "linux-build" { t.Fatalf("alias = %q, want linux-build", record.Alias) } if got := record.CommandTypes; len(got) != 2 || got[0] != "Shell" || got[1] != "Git" { t.Fatalf("command types = %#v, want [Shell Git]", got) } } func TestHandleRunnerRegisterRejectsMissingToken(t *testing.T) { registry := runnerregistry.New() body := bytes.NewBufferString(`{ "runner_id":"runner-123", "protocol_version":"oto.runner.v1" }`) req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body) rr := httptest.NewRecorder() handleRunnerRegister(registry)(rr, req) if rr.Code != http.StatusOK { t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK) } var response otopb.RegisterRunnerResponse if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil { t.Fatalf("decode response: %v", err) } if response.GetAccepted() { t.Fatal("missing token registration was accepted") } if response.GetRejectReason() != "missing enrollment token" { t.Fatalf("reject reason = %q, want missing enrollment token", response.GetRejectReason()) } } func TestHandleRunnerRegisterMethodNotAllowed(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/register", nil) rr := httptest.NewRecorder() handleRunnerRegister(runnerregistry.New())(rr, req) if rr.Code != http.StatusMethodNotAllowed { t.Fatalf("status = %v, want %v", rr.Code, http.StatusMethodNotAllowed) } } func TestHandleRunnerHeartbeat(t *testing.T) { registry := runnerregistry.New() // 1. Heartbeat on unregistered runner should fail (404) body := bytes.NewBufferString(`{"runner_id":"runner-123", "status":1}`) req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/heartbeat", body) rr := httptest.NewRecorder() handleRunnerHeartbeat(registry)(rr, req) if rr.Code != http.StatusNotFound { t.Fatalf("status = %v, want %v for unregistered runner", rr.Code, http.StatusNotFound) } // 2. Register runner first registry.Register(&otopb.RegisterRunnerRequest{ EnrollmentToken: "token-123", RunnerId: "runner-123", ProtocolVersion: "oto.runner.v1", }) // 3. Valid Heartbeat on registered runner should succeed body = bytes.NewBufferString(`{"runner_id":"runner-123", "status":1}`) req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/heartbeat", body) rr = httptest.NewRecorder() handleRunnerHeartbeat(registry)(rr, req) if rr.Code != http.StatusOK { t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK) } var response otopb.HeartbeatResponse if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil { t.Fatalf("decode heartbeat response: %v", err) } if !response.GetSuccess() { t.Fatal("heartbeat failed") } // Verify status is online record, _ := registry.Snapshot("runner-123") if record.Status != runnerregistry.StatusOnline { t.Fatalf("status = %q, want %q", record.Status, runnerregistry.StatusOnline) } } func TestHandleRunnerDisconnect(t *testing.T) { registry := runnerregistry.New() // 1. Disconnect unregistered runner should fail (404) req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/disconnect", nil) rr := httptest.NewRecorder() handleRunnerDisconnect(registry)(rr, req) if rr.Code != http.StatusNotFound { t.Fatalf("status = %v, want %v", rr.Code, http.StatusNotFound) } // 2. Register runner registry.Register(&otopb.RegisterRunnerRequest{ EnrollmentToken: "token-123", RunnerId: "runner-123", ProtocolVersion: "oto.runner.v1", }) // 3. Disconnect registered runner should succeed req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/disconnect", nil) rr = httptest.NewRecorder() handleRunnerDisconnect(registry)(rr, req) if rr.Code != http.StatusOK { t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK) } var res struct { Success bool `json:"success"` } if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil { t.Fatalf("decode response: %v", err) } if !res.Success { t.Fatal("expected success=true") } record, _ := registry.Snapshot("runner-123") if record.Status != runnerregistry.StatusDisconnected { t.Fatalf("status = %q, want %q", record.Status, runnerregistry.StatusDisconnected) } } func TestHandleGetRunner(t *testing.T) { registry := runnerregistry.New() // 1. Get unregistered runner should return 404 req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-123", nil) rr := httptest.NewRecorder() handleGetRunner(registry)(rr, req) if rr.Code != http.StatusNotFound { t.Fatalf("status = %v, want %v", rr.Code, http.StatusNotFound) } // 2. Register runner registry.Register(&otopb.RegisterRunnerRequest{ EnrollmentToken: "token-123", RunnerId: "runner-123", Alias: "build-linux", ProtocolVersion: "oto.runner.v1", }) // 3. Get registered runner should succeed (200) and return JSON record status req = httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-123", nil) rr = httptest.NewRecorder() handleGetRunner(registry)(rr, req) if rr.Code != http.StatusOK { t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK) } var res map[string]interface{} if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil { t.Fatalf("decode response: %v", err) } if res["runner_id"] != "runner-123" { t.Fatalf("runner_id = %v", res["runner_id"]) } if res["status"] != runnerregistry.StatusAccepted { t.Fatalf("status = %v", res["status"]) } } func TestHandleRunnerHeartbeatMismatch(t *testing.T) { registry := runnerregistry.New() // Register runner-a registry.Register(&otopb.RegisterRunnerRequest{ EnrollmentToken: "token-123", RunnerId: "runner-a", ProtocolVersion: "oto.runner.v1", }) // Register runner-b registry.Register(&otopb.RegisterRunnerRequest{ EnrollmentToken: "token-123", RunnerId: "runner-b", ProtocolVersion: "oto.runner.v1", }) // Heartbeat with mismatched runner ID in body vs path body := bytes.NewBufferString(`{"runner_id":"runner-b", "status":1}`) req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-a/heartbeat", body) rr := httptest.NewRecorder() handleRunnerHeartbeat(registry)(rr, req) if rr.Code != http.StatusBadRequest { t.Fatalf("status = %v, want %v", rr.Code, http.StatusBadRequest) } var response otopb.HeartbeatResponse if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil { t.Fatalf("decode response: %v", err) } if response.GetSuccess() { t.Fatal("expected heartbeat to fail due to ID mismatch") } if response.GetErrorMessage() != "runner id mismatch between path and body" { t.Fatalf("error message = %q", response.GetErrorMessage()) } // Verify runner-b remains accepted (not revived/online) rec, _ := registry.Snapshot("runner-b") if rec.Status != runnerregistry.StatusAccepted { t.Fatalf("runner-b status = %q, want %q", rec.Status, runnerregistry.StatusAccepted) } }