Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158package helps
import ( "context" "errors" "net/http" "net/http/httptest" "strings" "testing"
"github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor")
func TestRequestLoggingDoesNotMarkUpstreamAttempt(t *testing.T) { tests := []struct { name string record func(context.Context) }{ { name: "HTTP", record: func(ctx context.Context) { RecordAPIRequest(ctx, &config.Config{}, UpstreamRequestLog{URL: "https://api.example.com", Method: http.MethodPost}) }, }, { name: "websocket", record: func(ctx context.Context) { RecordAPIWebsocketRequest(ctx, &config.Config{}, UpstreamRequestLog{URL: "wss://api.example.com", Method: "WEBSOCKET"}) }, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) test.record(ctx) if cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("request logging marked an upstream attempt before transport") } }) }}
func TestRecordAPIRequestClonesDeferredBodyWhenRequestLogDisabled(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(recorder) ctx := context.WithValue(context.Background(), "gin", ginCtx) body := []byte(`{"model":"original"}`)
RecordAPIRequest(ctx, &config.Config{}, UpstreamRequestLog{ URL: "https://api.example.com/v1/responses", Method: http.MethodPost, Body: body, }) body[10] = 'X'
value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey) if !exists { t.Fatal("deferred API request was not captured") } requests, ok := value.([]logging.DeferredAPIRequest) if !ok || len(requests) != 1 { t.Fatalf("deferred API requests = %#v, want one request", value) } captured := string(requests[0]()) if !strings.Contains(captured, `{"model":"original"}`) { t.Fatalf("captured API request = %q, want original body", captured) }}
func TestRecordAPIResponseMetadataStoresHeadersWhenRequestLogDisabled(t *testing.T) { ctx := logging.WithResponseHeadersHolder(context.Background()) headers := http.Header{} headers.Add("X-Upstream-Request-Id", "upstream-req-1")
RecordAPIResponseMetadata(ctx, &config.Config{}, http.StatusOK, headers) headers.Set("X-Upstream-Request-Id", "mutated")
got := logging.GetResponseHeaders(ctx) if got.Get("X-Upstream-Request-Id") != "upstream-req-1" { t.Fatalf("response header = %q, want %q", got.Get("X-Upstream-Request-Id"), "upstream-req-1") }}
func TestAPIResponseAttemptsAreSeparated(t *testing.T) { gin.SetMode(gin.TestMode)
tests := []struct { name string fileBacked bool firstResponseBody []byte }{ {name: "memory backed error"}, {name: "file backed error", fileBacked: true}, {name: "memory backed partial body", firstResponseBody: []byte("partial")}, {name: "file backed partial body", fileBacked: true, firstResponseBody: []byte("partial")}, }
for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { recorder := httptest.NewRecorder() ginCtx, _ := gin.CreateTestContext(recorder) var responseSource *logging.FileBodySource if tt.fileBacked { var errSource error responseSource, errSource = logging.NewFileBodySourceInDir(t.TempDir(), "api-response") if errSource != nil { t.Fatalf("NewFileBodySourceInDir: %v", errSource) } t.Cleanup(func() { if errCleanup := responseSource.Cleanup(); errCleanup != nil { t.Errorf("Cleanup: %v", errCleanup) } }) ginCtx.Set(logging.APIResponseSourceContextKey, responseSource) }
ctx := context.WithValue(context.Background(), "gin", ginCtx) cfg := &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}} RecordAPIRequest(ctx, cfg, UpstreamRequestLog{URL: "https://api.example.com/first", Method: http.MethodPost}) if len(tt.firstResponseBody) > 0 { AppendAPIResponseChunk(ctx, cfg, tt.firstResponseBody) } else { RecordAPIResponseError(ctx, cfg, errors.New("EOF")) } RecordAPIRequest(ctx, cfg, UpstreamRequestLog{URL: "https://api.example.com/second", Method: http.MethodPost}) RecordAPIResponseError(ctx, cfg, errors.New("retry failed"))
var response []byte if responseSource != nil { var errBytes error response, errBytes = responseSource.Bytes() if errBytes != nil { t.Fatalf("responseSource.Bytes: %v", errBytes) } } else { value, exists := ginCtx.Get(apiResponseKey) if !exists { t.Fatal("API_RESPONSE was not captured") } response, _ = value.([]byte) }
previousEnd := "Error: EOF" if len(tt.firstResponseBody) > 0 { previousEnd = string(tt.firstResponseBody) } wantBoundary := previousEnd + "\n\n=== API RESPONSE 2 ===" if !strings.Contains(string(response), wantBoundary) { t.Fatalf("API response attempts are not separated by one blank line:\n%q\nwant boundary %q", response, wantBoundary) } }) }}