zero/internal/aimlapi/stream_test.go

package aimlapi

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestStreamTopUpByKeyFundsWithoutExchange(t *testing.T) {
	var topupCalls, exchangeCalls int
	var topupBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/partner-checkout/sessions"):
			_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pcs_tok", Status: SessionStatusPendingAuth})
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/v2/billing/topup"):
			topupCalls++
			_ = json.NewDecoder(r.Body).Decode(&topupBody)
			_ = json.NewEncoder(w).Encode(TopUpByKeyResult{Checkout: PaymentSession{PayURL: "https://pay/checkout"}})
		case r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/partner-checkout/sessions/"):
			_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pcs_tok", Status: SessionStatusPaid})
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/exchange"):
			exchangeCalls++
			t.Errorf("by-key must never exchange a key")
		default:
			t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)
	t.Setenv("AIMLAPI_INFERENCE_URL", server.URL+"/v1")

	var opened string
	result, err := StreamTopUpByKey(context.Background(), StreamTopUpByKeyOptions{
		APIKey:           "sk-user",
		AmountUSD:        "20",
		AutoTopUp:        true,
		PaymentSessionID: "pay-1",
		OpenBrowser:      func(u string) error { opened = u; return nil },
	})
	if err != nil {
		t.Fatalf("StreamTopUpByKey = %v", err)
	}
	if result.APIKey != "" {
		t.Fatalf("by-key must not mint a key, got %q", result.APIKey)
	}
	if topupCalls != 1 || exchangeCalls != 0 {
		t.Fatalf("topup=%d exchange=%d, want 1/0", topupCalls, exchangeCalls)
	}
	if opened != "https://pay/checkout" {
		t.Fatalf("opened checkout = %q", opened)
	}
	if topupBody["autoTopUp"] != true || topupBody["paymentSessionId"] != "pay-1" {
		t.Fatalf("topup body missing autoTopUp/paymentSessionId: %#v", topupBody)
	}
}

func TestValidatedCheckoutURLRequiresAbsoluteHTTPS(t *testing.T) {
	tests := []struct {
		name    string
		value   string
		want    string
		wantErr bool
	}{
		{name: "https", value: " https://pay.example.test/checkout?id=1 ", want: "https://pay.example.test/checkout?id=1"},
		{name: "empty", wantErr: true},
		{name: "relative", value: "/checkout", wantErr: true},
		{name: "http", value: "http://pay.example.test/checkout", wantErr: true},
		{name: "file", value: "file:///tmp/checkout", wantErr: true},
		{name: "custom protocol", value: "wallet://checkout/123", wantErr: true},
		{name: "userinfo", value: "https://[email protected]/checkout", wantErr: true},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			got, err := validatedCheckoutURL(test.value)
			if (err != nil) != test.wantErr {
				t.Fatalf("validatedCheckoutURL(%q) error = %v", test.value, err)
			}
			if got != test.want {
				t.Fatalf("validatedCheckoutURL(%q) = %q, want %q", test.value, got, test.want)
			}
		})
	}
}

func TestStreamTopUpByKeyPinsValidatedInferenceEndpoint(t *testing.T) {
	validatedCalls := 0
	validated := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/partner-checkout/sessions"):
			_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pcs_tok", Status: SessionStatusPendingAuth})
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/v2/billing/topup"):
			validatedCalls++
			_ = json.NewEncoder(w).Encode(TopUpByKeyResult{Checkout: PaymentSession{PayURL: "https://pay/checkout"}})
		case r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/partner-checkout/sessions/"):
			_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pcs_tok", Status: SessionStatusPaid})
		default:
			t.Errorf("unexpected validated-endpoint request: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer validated.Close()
	override := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Errorf("validated key leaked to process override: %s %s", r.Method, r.URL.Path)
	}))
	defer override.Close()
	t.Setenv("AIMLAPI_APP_URL", validated.URL)
	t.Setenv("AIMLAPI_INFERENCE_URL", override.URL)

	_, err := StreamTopUpByKey(context.Background(), StreamTopUpByKeyOptions{
		APIKey:           "production-key",
		AmountUSD:        "20",
		InferenceBaseURL: validated.URL,
		PaymentSessionID: "payment-1",
		OpenBrowser:      func(string) error { return nil },
	})
	if err != nil {
		t.Fatalf("StreamTopUpByKey() error = %v", err)
	}
	if validatedCalls != 1 {
		t.Fatalf("validated endpoint top-up calls = %d, want 1", validatedCalls)
	}
}

func TestStreamTopUpEmailRetrySendsStablePaymentSessionID(t *testing.T) {
	getCalls := 0
	var paymentSessionID string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/partner-checkout/sessions/"):
			status := SessionStatusPendingAuth
			if getCalls > 0 {
				status = SessionStatusPaid
			}
			getCalls++
			_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pcs_retry", Status: status})
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/pay"):
			var body map[string]any
			_ = json.NewDecoder(r.Body).Decode(&body)
			paymentSessionID, _ = body["paymentSessionId"].(string)
			_ = json.NewEncoder(w).Encode(PayResult{Checkout: PaymentSession{PayURL: "https://pay/checkout"}})
		default:
			t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	_, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pcs_retry",
		PaymentSessionID:   "payment-stable",
		AmountUSD:          "20",
		NoOpen:             true,
	})
	if err != nil {
		t.Fatalf("StreamTopUp() error = %v", err)
	}
	if paymentSessionID != "payment-stable" {
		t.Fatalf("paymentSessionId = %q, want stable retry id", paymentSessionID)
	}
}

// resumeServer is a partner-checkout backend fake that fails the test if the
// re-charging endpoints (create-session / pay) are ever hit, and serves a
// scripted GetSession status sequence so a resume can be driven deterministically.
type resumeServer struct {
	t          *testing.T
	statuses   []string // GetSession status per call, last value repeats
	getCalls   int
	payCalls   int
	createCall int
	exchanged  int
}

func (rs *resumeServer) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/partner-checkout/sessions"):
			rs.createCall++
			rs.t.Errorf("resume must not create a new session")
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/pay"):
			rs.payCalls++
			rs.t.Errorf("resume must not re-pay (double charge)")
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/exchange"):
			rs.exchanged++
			_ = json.NewEncoder(w).Encode(ExchangeResult{APIKey: "sk-resumed", APIKeyID: "key_1"})
		case r.Method == http.MethodGet:
			status := rs.statuses[len(rs.statuses)-1]
			if rs.getCalls < len(rs.statuses) {
				status = rs.statuses[rs.getCalls]
			}
			rs.getCalls++
			_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pc_tok", Status: SessionStatus(status)})
		default:
			rs.t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path)
		}
	})
}

func TestStreamTopUpResumePaidExchangesWithoutRecharge(t *testing.T) {
	rs := &resumeServer{t: t, statuses: []string{"paid"}}
	server := httptest.NewServer(rs.handler())
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	var seen string
	result, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pc_tok",
		PaymentSessionID:   "payment-1",
		AmountUSD:          "20",
		Exchange:           true,
		NoOpen:             true,
		OnSession:          func(token string) { seen = token },
	})
	if err != nil {
		t.Fatalf("StreamTopUp resume = %v", err)
	}
	if result.APIKey != "sk-resumed" {
		t.Fatalf("APIKey = %q, want recovered key", result.APIKey)
	}
	if rs.exchanged != 1 {
		t.Fatalf("exchange calls = %d, want 1", rs.exchanged)
	}
	if seen != "pc_tok" {
		t.Fatalf("OnSession token = %q, want retained pc_tok", seen)
	}
}

func TestStreamTopUpResumePendingPaymentPollsNeverRepays(t *testing.T) {
	// First GetSession (resume resolution) sees the still-open checkout; the next
	// (poll) sees it settle. Pay must never be called — the fake fails the test if
	// it is.
	rs := &resumeServer{t: t, statuses: []string{"pending_payment", "paid"}}
	server := httptest.NewServer(rs.handler())
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	result, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pc_tok",
		PaymentSessionID:   "payment-1",
		AmountUSD:          "20",
		Exchange:           true,
		NoOpen:             true,
	})
	if err != nil {
		t.Fatalf("StreamTopUp resume = %v", err)
	}
	if result.APIKey != "sk-resumed" {
		t.Fatalf("APIKey = %q, want recovered key", result.APIKey)
	}
	if rs.payCalls != 0 {
		t.Fatalf("pay calls = %d, want 0", rs.payCalls)
	}
}

func TestStreamTopUpPollingDeadSessionClearsRetainedToken(t *testing.T) {
	rs := &resumeServer{t: t, statuses: []string{"pending_payment", "expired"}}
	server := httptest.NewServer(rs.handler())
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	seen := []string{}
	_, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pc_tok",
		PaymentSessionID:   "payment-1",
		AmountUSD:          "20",
		Exchange:           true,
		NoOpen:             true,
		OnSession:          func(token string) { seen = append(seen, token) },
	})
	if err == nil || !strings.Contains(err.Error(), "payment expired") {
		t.Fatalf("err = %v, want expired payment error", err)
	}
	if len(seen) == 0 || seen[len(seen)-1] != "" {
		t.Fatalf("OnSession events = %#v, want terminal clear", seen)
	}
}

func TestPollUntilPaidTerminal4xxClearsRetainedToken(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "session rejected", http.StatusNotFound)
	}))
	defer server.Close()

	seen := "retained"
	_, err := pollUntilPaid(context.Background(), NewClient(Endpoints{AppBaseURL: server.URL}, server.Client()), "pc_missing", func(token string) {
		seen = token
	})
	if err == nil {
		t.Fatal("pollUntilPaid() error = nil, want terminal 4xx")
	}
	if seen != "" {
		t.Fatalf("retained session = %q, want cleared", seen)
	}
}

func TestPollUntilExchangeSettledClearsDeadSession(t *testing.T) {
	for _, test := range []struct {
		name       string
		statusCode int
		status     SessionStatus
	}{
		{name: "terminal API rejection", statusCode: http.StatusNotFound},
		{name: "expired exchange claim", status: SessionStatusExpired},
	} {
		t.Run(test.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				if test.statusCode != 0 {
					http.Error(w, "session rejected", test.statusCode)
					return
				}
				_ = json.NewEncoder(w).Encode(PartnerCheckoutSession{SessionToken: "pcs_dead", Status: test.status})
			}))
			defer server.Close()

			seen := "retained"
			err := pollUntilExchangeSettled(context.Background(), NewClient(Endpoints{AppBaseURL: server.URL}, server.Client()), "pcs_dead", func(token string) {
				seen = token
			})
			if err == nil {
				t.Fatal("pollUntilExchangeSettled() error = nil")
			}
			if seen != "" {
				t.Fatalf("retained session = %q, want cleared", seen)
			}
		})
	}
}

func TestResumeHelpersClearOnlyRejectedSessions(t *testing.T) {
	tests := []struct {
		name      string
		status    int
		wantClear bool
	}{
		{name: "terminal 4xx", status: http.StatusNotFound, wantClear: true},
		{name: "retryable 5xx", status: http.StatusServiceUnavailable, wantClear: false},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				http.Error(w, "session lookup failed", test.status)
			}))
			defer server.Close()
			client := NewClient(Endpoints{AppBaseURL: server.URL}, server.Client())

			for _, flow := range []struct {
				name    string
				resolve func(func(string)) error
			}{
				{
					name: "session bearer",
					resolve: func(onSession func(string)) error {
						_, _, err := resolveTopupSession(context.Background(), client, StreamTopUpOptions{
							ResumeSessionToken: "pc_rejected", OnSession: onSession,
						}, DefaultPartnerID, DefaultPartnerName, Endpoints{})
						return err
					},
				},
				{
					name: "api key",
					resolve: func(onSession func(string)) error {
						_, _, err := resolveByKeySession(context.Background(), client, StreamTopUpByKeyOptions{
							ResumeSessionToken: "pc_rejected", OnSession: onSession,
						}, DefaultPartnerID, DefaultPartnerName, Endpoints{})
						return err
					},
				},
			} {
				t.Run(flow.name, func(t *testing.T) {
					seen := "retained"
					if err := flow.resolve(func(token string) { seen = token }); err == nil {
						t.Fatal("resolve error = nil")
					}
					if cleared := seen == ""; cleared != test.wantClear {
						t.Fatalf("session cleared = %v, want %v", cleared, test.wantClear)
					}
				})
			}
		})
	}
}

func TestStreamTopUpResumeExchangedReturnsRecoveryError(t *testing.T) {
	rs := &resumeServer{t: t, statuses: []string{"exchanged"}}
	server := httptest.NewServer(rs.handler())
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	_, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pc_tok",
		PaymentSessionID:   "payment-1",
		AmountUSD:          "20",
		Exchange:           true,
		NoOpen:             true,
	})
	if err == nil || !strings.Contains(err.Error(), "already exchanged") {
		t.Fatalf("err = %v, want an already-exchanged recovery error", err)
	}
	if rs.exchanged != 0 {
		t.Fatalf("exchange calls = %d, want 0 (no second key)", rs.exchanged)
	}
}

func TestStreamTopUpResumeExchangingPollsWithoutSecondExchange(t *testing.T) {
	rs := &resumeServer{t: t, statuses: []string{"exchanging", "exchanged"}}
	server := httptest.NewServer(rs.handler())
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	_, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pc_tok",
		PaymentSessionID:   "payment-1",
		AmountUSD:          "20",
		Exchange:           true,
		NoOpen:             true,
	})
	if err == nil || !strings.Contains(err.Error(), "already exchanged") {
		t.Fatalf("err = %v, want dashboard recovery guidance", err)
	}
	if rs.exchanged != 0 {
		t.Fatalf("exchange calls = %d, want 0 while an exchange claim is in flight", rs.exchanged)
	}
	if rs.createCall != 0 || rs.payCalls != 0 {
		t.Fatalf("resume created/paid a new checkout: create=%d pay=%d", rs.createCall, rs.payCalls)
	}
}

func TestStreamTopUpPollObservesExchangingWithoutSecondExchange(t *testing.T) {
	rs := &resumeServer{t: t, statuses: []string{"pending_payment", "exchanging", "exchanged"}}
	server := httptest.NewServer(rs.handler())
	defer server.Close()
	t.Setenv("AIMLAPI_APP_URL", server.URL)

	_, err := StreamTopUp(context.Background(), StreamTopUpOptions{
		SessionToken:       "bearer",
		ResumeSessionToken: "pc_tok",
		PaymentSessionID:   "payment-1",
		AmountUSD:          "20",
		Exchange:           true,
		NoOpen:             true,
	})
	if err == nil || !strings.Contains(err.Error(), "already exchanged") {
		t.Fatalf("err = %v, want dashboard recovery guidance", err)
	}
	if rs.exchanged != 0 {
		t.Fatalf("exchange calls = %d, want 0 after poll observed an existing claim", rs.exchanged)
	}
}