promptkit/internal/llm/openrouter_authentication_test.go

package llm

import (
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
	"testing"

	"go.maximumdirect.net/promptkit/internal/domain"
)

func TestProbeOpenRouterKey(t *testing.T) {
	valid := `{"data":{"label":"x","is_free_tier":false,"usage":0}}`
	for _, tc := range []struct {
		name                   string
		status                 int
		body                   string
		key                    string
		accepted               *bool
		unrecognized, overflow bool
	}{
		{name: "recognized", status: 200, body: valid, key: "secret", accepted: boolValue(true)},
		{name: "exact limit", status: 200, body: valid + strings.Repeat(" ", maxAuthenticationResponseBytes-len(valid)), accepted: boolValue(true)},
		{name: "unknown", status: 200, body: `{"data":{"label":"x","is_free_tier":false,"usage":"0"}}`, unrecognized: true},
		{name: "trailing value", status: 200, body: valid + ` {}`, unrecognized: true},
		{name: "overflow", status: 200, body: strings.Repeat("x", maxAuthenticationResponseBytes+1), overflow: true},
		{name: "unauthorized anonymous", status: 401, accepted: boolValue(false)},
		{name: "other status", status: 503},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var authorization string
			client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
				if req.URL.User != nil {
					t.Fatal("request retained URL credentials")
				}
				authorization = req.Header.Get("Authorization")
				return &http.Response{StatusCode: tc.status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(tc.body))}, nil
			})}})
			if err != nil {
				t.Fatal(err)
			}
			got, err := client.ProbeOpenRouterKey(context.Background(), "https://user:password@example.invalid/key", 0, domain.ExecutionTarget{APIKey: tc.key})
			if err != nil {
				t.Fatal(err)
			}
			if got.HTTP.Endpoint != "https://example.invalid/key" || got.HTTP.HTTPStatus == nil || *got.HTTP.HTTPStatus != tc.status || got.Unrecognized != tc.unrecognized || got.Overflow != tc.overflow {
				t.Fatalf("evidence = %+v", got)
			}
			if (got.Accepted == nil) != (tc.accepted == nil) || got.Accepted != nil && *got.Accepted != *tc.accepted {
				t.Fatalf("acceptance = %v", got.Accepted)
			}
			if (authorization != "") != (tc.key != "") || got.HTTP.CredentialsSupplied != (tc.key != "") {
				t.Fatalf("credential evidence = %+v", got.HTTP)
			}
		})
	}
}

func boolValue(value bool) *bool { return &value }

type authenticationFailingBody struct{ err error }

func (b authenticationFailingBody) Read([]byte) (int, error) { return 0, b.err }
func (b authenticationFailingBody) Close() error             { return nil }

func TestProbeOpenRouterKeyRetainsBodyFailure(t *testing.T) {
	want := errors.New("interrupted")
	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: 200, Header: make(http.Header), Body: authenticationFailingBody{err: want}}, nil
	})}})
	if err != nil {
		t.Fatal(err)
	}
	got, err := client.ProbeOpenRouterKey(context.Background(), "https://example.invalid/key", 0, domain.ExecutionTarget{})
	if !errors.Is(err, want) || !got.HTTP.ResponseReceived {
		t.Fatalf("evidence = %+v, error = %v", got, err)
	}
}

type observedAuthenticationBody struct {
	reader io.Reader
	read   int
	closed bool
}

func (b *observedAuthenticationBody) Read(p []byte) (int, error) {
	n, err := b.reader.Read(p)
	b.read += n
	return n, err
}

func (b *observedAuthenticationBody) Close() error {
	b.closed = true
	return nil
}

func TestProbeOpenRouterKeyBoundsAndClosesBodies(t *testing.T) {
	for _, tc := range []struct {
		name     string
		status   int
		wantRead int
	}{
		{"overflow", http.StatusOK, maxAuthenticationResponseBytes + 1},
		{"unauthorized", http.StatusUnauthorized, 0},
		{"other status", http.StatusServiceUnavailable, 0},
	} {
		t.Run(tc.name, func(t *testing.T) {
			body := &observedAuthenticationBody{reader: strings.NewReader(strings.Repeat("x", 2*maxAuthenticationResponseBytes))}
			client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
				return &http.Response{StatusCode: tc.status, ContentLength: 1, Header: make(http.Header), Body: body}, nil
			})}})
			if err != nil {
				t.Fatal(err)
			}
			got, err := client.ProbeOpenRouterKey(context.Background(), "https://example.invalid/key", 0, domain.ExecutionTarget{})
			if err != nil || body.read != tc.wantRead || !body.closed || got.Overflow != (tc.status == http.StatusOK) {
				t.Fatalf("evidence = %+v, read = %d, closed = %v, error = %v", got, body.read, body.closed, err)
			}
		})
	}
}

func TestOpenRouterKeyEndpoint(t *testing.T) {
	got, err := OpenRouterKeyEndpoint("https://example.invalid/api/v1/")
	if err != nil || got != "https://example.invalid/api/v1/key" {
		t.Fatalf("endpoint = %q, error = %v", got, err)
	}
}

func TestProbeOpenRouterKeyPreservesCancellation(t *testing.T) {
	client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
		<-req.Context().Done()
		return nil, req.Context().Err()
	})}})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	got, err := client.ProbeOpenRouterKey(ctx, "https://example.invalid/key", 0, domain.ExecutionTarget{})
	if !errors.Is(err, context.Canceled) || got.Accepted != nil || got.HTTP.ResponseReceived {
		t.Fatalf("evidence = %+v, error = %v", got, err)
	}
}