knowledge-governance/internal/qaopsai/client_test.go

package qaopsai

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

func TestClientUsesConfiguredModelAndWriteOnlyKeyWithoutChangingInput(t *testing.T) {
	count := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		count++
		if r.URL.Path != "/v1/chat/completions" || r.Header.Get("Authorization") != "Bearer fixture-only" {
			t.Error("wrong route or authorization")
		}
		var request struct {
			Model          string                           `json:"model"`
			Messages       []struct{ Role, Content string } `json:"messages"`
			ResponseFormat map[string]string                `json:"response_format"`
		}
		if json.NewDecoder(r.Body).Decode(&request) != nil || request.Model != "fixture-model" || len(request.Messages) != 2 || request.ResponseFormat["type"] != "json_object" {
			t.Error("configured protocol not used")
		}
		if len(request.Messages) == 2 && strings.Contains(request.Messages[1].Content, "fixture-only") {
			t.Error("key copied into business prompt")
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"model":"fixture-returned-model","choices":[{"message":{"content":"{\"ok\":true}"},"finish_reason":"stop"}]}`))
	}))
	defer server.Close()
	got, err := (Client{AllowPrivate: true}).Call(context.Background(), server.URL+"/v1", "fixture-model", "fixture-only", "Respond only with JSON", map[string]string{"question": "synthetic only"}, 64)
	if err != nil || got.Content != `{"ok":true}` || got.Model != "fixture-returned-model" || count != 1 {
		t.Fatalf("actual compatibility request not completed: response=%+v error=%v calls=%d", got, err, count)
	}
}

func TestClientReleasesPerCallTransportConnections(t *testing.T) {
	closed := make(chan struct{}, 1)
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`{"model":"fixture","choices":[{"message":{"content":"{\"ok\":true}"},"finish_reason":"stop"}]}`))
	}))
	server.Config.ConnState = func(_ net.Conn, state http.ConnState) {
		if state == http.StateClosed {
			select {
			case closed <- struct{}{}:
			default:
			}
		}
	}
	server.Start()
	defer server.Close()
	if _, err := (Client{AllowPrivate: true}).Call(context.Background(), server.URL, "fixture", "fixture-only", "json", map[string]string{}, 64); err != nil {
		t.Fatal(err)
	}
	select {
	case <-closed:
	case <-time.After(time.Second):
		t.Fatal("completed per-call transport left an idle model connection open")
	}
}

func TestClientRejectsPrivateEndpointsBeforeSendingSecrets(t *testing.T) {
	for _, encryptedDNS := range []bool{false, true} {
		for _, base := range []string{"https://127.0.0.1:1/v1", "https://198.18.0.40/v1", "https://100.92.49.78/v1", "https://10.0.0.1/v1"} {
			_, err := (Client{DNSOverTLS: encryptedDNS}).Call(context.Background(), base, "fixture", "fixture-only", "json", map[string]string{}, 32)
			if err != ErrInvalid {
				t.Fatalf("private/reserved destination accepted with encryptedDNS=%v: %v", encryptedDNS, err)
			}
		}
	}
}

func TestClientDistinguishesTruncatedOutputAndRetainsProviderCooldown(t *testing.T) {
	for _, tc := range []struct {
		name, body, code string
		status           int
		delay            time.Duration
	}{
		{"truncated", `{"model":"fixture","choices":[{"message":{"content":"{"},"finish_reason":"length"}]}`, "output_truncated", 200, 0},
		{"rate-limit", `{}`, "rate_limited", 429, 45 * time.Second},
	} {
		t.Run(tc.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Retry-After", "45")
				w.WriteHeader(tc.status)
				_, _ = w.Write([]byte(tc.body))
			}))
			defer server.Close()
			_, err := (Client{AllowPrivate: true}).Call(context.Background(), server.URL, "fixture", "fixture-only", "json", map[string]string{}, 64)
			var remote *RemoteError
			if !errors.As(err, &remote) || remote.Code != tc.code || !remote.Retryable || remote.RetryAfter != tc.delay {
				t.Fatalf("lost retry classification/cooldown: %+v", err)
			}
		})
	}
}