1 files changed, 89 insertions, 0 deletions
diff --git a/internal/operations/operations_test.go b/internal/operations/operations_test.go
new file mode 100644
index 0000000..6936b1a
--- /dev/null
+++ b/internal/operations/operations_test.go
@@ -0,0 +1,89 @@
+package operations
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"aigw/internal/billing"
+	"aigw/internal/controlplane"
+)
+
+type healthyStore struct{}
+
+func (healthyStore) Ping(context.Context) error { return nil }
+func (healthyStore) MailQueueStatus(context.Context) (controlplane.MailQueueStatus, error) {
+	return controlplane.MailQueueStatus{}, nil
+}
+
+type slowBilling struct{}
+
+func (slowBilling) Ping(context.Context) error { return nil }
+func (slowBilling) SettlementQueueStatus(context.Context) (billing.SettlementQueueStatus, error) {
+	return billing.SettlementQueueStatus{}, nil
+}
+func (slowBilling) OperationalStatus(ctx context.Context) (billing.OperationalStatus, error) {
+	<-ctx.Done()
+	return billing.OperationalStatus{}, ctx.Err()
+}
+
+func TestReadinessTimeoutDoesNotMislabelCompletedChecks(t *testing.T) {
+	handler := Handler{Store: healthyStore{}, Billing: slowBilling{}, ReadinessTimeout: 20 * time.Millisecond}
+	request := httptest.NewRequest(http.MethodGet, "/readyz", nil)
+	recorder := httptest.NewRecorder()
+	handler.ServeHTTP(recorder, request)
+	if recorder.Code != http.StatusServiceUnavailable {
+		t.Fatalf("status = %d, want 503", recorder.Code)
+	}
+	var payload struct {
+		Checks map[string]struct {
+			Status string `json:"status"`
+			Error  string `json:"error"`
+		} `json:"checks"`
+	}
+	if err := json.NewDecoder(recorder.Body).Decode(&payload); err != nil {
+		t.Fatal(err)
+	}
+	for _, name := range []string{"postgres", "mail_queue", "billing_postgres", "settlement_queue"} {
+		if payload.Checks[name].Status != "ok" {
+			t.Fatalf("%s status = %q, want ok; payload=%s", name, payload.Checks[name].Status, recorder.Body.String())
+		}
+	}
+	if payload.Checks["billing_operations"].Status != "failed" || payload.Checks["billing_operations"].Error != "check timed out" {
+		t.Fatalf("billing operations = %+v", payload.Checks["billing_operations"])
+	}
+}
+
+func TestMailQueueIsCheckedWithoutBilling(t *testing.T) {
+	handler := Handler{Store: healthyStore{}, ReadinessTimeout: time.Second}
+	recorder := httptest.NewRecorder()
+	handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/readyz", nil))
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String())
+	}
+	var payload struct {
+		Checks map[string]any `json:"checks"`
+	}
+	if err := json.NewDecoder(recorder.Body).Decode(&payload); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := payload.Checks["mail_queue"]; !ok {
+		t.Fatal("mail_queue check missing when billing is disabled")
+	}
+}
+
+func TestReadinessChecksDoNotInheritCanceledClientContext(t *testing.T) {
+	handler := Handler{Store: healthyStore{}, ReadinessTimeout: time.Second}
+	request := httptest.NewRequest(http.MethodGet, "/readyz", nil)
+	ctx, cancel := context.WithCancel(request.Context())
+	cancel()
+	request = request.WithContext(ctx)
+	recorder := httptest.NewRecorder()
+	handler.ServeHTTP(recorder, request)
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String())
+	}
+}