package pluginhost import ( "bytes" "context" "encoding/json" "net/http" "net/http/httptest" "testing" "time" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) func TestHostHTTPDoCallbackUsesHostHTTPClient(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { t.Fatalf("method = %s, want POST", r.Method) } w.Header().Set("X-Test", "ok") _, _ = w.Write([]byte(`{"ok":true}`)) })) defer server.Close() req := pluginapi.HTTPRequest{ Method: http.MethodPost, URL: server.URL, Body: []byte(`{"request":true}`), } rawReq, errMarshal := json.Marshal(req) if errMarshal != nil { t.Fatalf("marshal request: %v", errMarshal) } rawResp, errCall := New().callFromPlugin(context.Background(), pluginabi.MethodHostHTTPDo, rawReq) if errCall != nil { t.Fatalf("callFromPlugin() error = %v", errCall) } resp, errDecode := decodeRPCEnvelope[pluginapi.HTTPResponse](rawResp) if errDecode != nil { t.Fatalf("decode response: %v", errDecode) } if resp.StatusCode != http.StatusOK || string(resp.Body) != `{"ok":true}` { t.Fatalf("response = %#v, want status 200 body", resp) } if resp.Headers.Get("X-Test") != "ok" { t.Fatalf("X-Test = %q, want ok", resp.Headers.Get("X-Test")) } } func TestHostHTTPDoCallbackRestoresRegisteredRequestContext(t *testing.T) { gin.SetMode(gin.TestMode) ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) ctx := context.WithValue(context.Background(), "gin", ginCtx) host := New() host.mu.Lock() host.runtimeConfig = &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}} host.mu.Unlock() callbackID, closeCallback := host.openCallbackContext(ctx) defer closeCallback() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Context().Err() != nil { t.Fatalf("request context error = %v", r.Context().Err()) } w.Header().Set("X-Upstream", "ok") _, _ = w.Write([]byte("upstream-body")) })) defer server.Close() rawReq, errMarshal := json.Marshal(rpcHostHTTPRequest{ HostCallbackID: callbackID, Method: http.MethodPost, URL: server.URL, Body: []byte(`{"request":true}`), }) if errMarshal != nil { t.Fatalf("marshal request: %v", errMarshal) } if _, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostHTTPDo, rawReq); errCall != nil { t.Fatalf("callFromPlugin() error = %v", errCall) } rawAPIRequest, okRequest := ginCtx.Get("API_REQUEST") if !okRequest { t.Fatal("API_REQUEST was not captured on the original Gin context") } apiRequest, _ := rawAPIRequest.([]byte) if !bytes.Contains(apiRequest, []byte("=== API REQUEST 1 ===")) || !bytes.Contains(apiRequest, []byte(`{"request":true}`)) { t.Fatalf("API_REQUEST = %q, want upstream request details", apiRequest) } rawAPIResponse, okResponse := ginCtx.Get("API_RESPONSE") if !okResponse { t.Fatal("API_RESPONSE was not captured on the original Gin context") } apiResponse, _ := rawAPIResponse.([]byte) if !bytes.Contains(apiResponse, []byte("=== API RESPONSE 1 ===")) || !bytes.Contains(apiResponse, []byte("upstream-body")) { t.Fatalf("API_RESPONSE = %q, want upstream response details", apiResponse) } } func TestHostHTTPDoStreamCallbackReturnsBeforeUpstreamCompletes(t *testing.T) { release := make(chan struct{}) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("first")) if flusher, ok := w.(http.Flusher); ok { flusher.Flush() } <-release _, _ = w.Write([]byte("second")) })) defer server.Close() defer close(release) rawReq, errMarshal := json.Marshal(pluginapi.HTTPRequest{ Method: http.MethodGet, URL: server.URL, }) if errMarshal != nil { t.Fatalf("marshal request: %v", errMarshal) } type callResult struct { raw []byte err error } done := make(chan callResult, 1) host := New() go func() { rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostHTTPDoStream, rawReq) done <- callResult{raw: rawResp, err: errCall} }() var result callResult select { case result = <-done: case <-time.After(time.Second): t.Fatal("host.http.do_stream waited for the whole upstream response") } if result.err != nil { t.Fatalf("callFromPlugin() error = %v", result.err) } resp, errDecode := decodeRPCEnvelope[rpcHostHTTPStreamResponse](result.raw) if errDecode != nil { t.Fatalf("decode response: %v", errDecode) } if resp.StreamID == "" { t.Fatalf("stream id is empty: %#v", resp) } readReq, errMarshal := json.Marshal(rpcHostHTTPStreamReadRequest{StreamID: resp.StreamID}) if errMarshal != nil { t.Fatalf("marshal read request: %v", errMarshal) } rawRead, errRead := host.callFromPlugin(context.Background(), pluginabi.MethodHostHTTPStreamRead, readReq) if errRead != nil { t.Fatalf("read callback error = %v", errRead) } chunk, errDecode := decodeRPCEnvelope[rpcHostHTTPStreamReadResponse](rawRead) if errDecode != nil { t.Fatalf("decode read response: %v", errDecode) } if string(chunk.Payload) != "first" || chunk.Done || chunk.Error != "" { t.Fatalf("read chunk = %#v, want first payload", chunk) } closeReq, errMarshal := json.Marshal(rpcHostHTTPStreamCloseRequest{StreamID: resp.StreamID}) if errMarshal != nil { t.Fatalf("marshal close request: %v", errMarshal) } if _, errClose := host.callFromPlugin(context.Background(), pluginabi.MethodHostHTTPStreamClose, closeReq); errClose != nil { t.Fatalf("close callback error = %v", errClose) } } func TestHostStreamCallbacksEmitAndClose(t *testing.T) { host := New() streamID, chunks, cleanup := host.streams.open(context.Background()) defer cleanup() emitReq, errMarshal := json.Marshal(rpcStreamEmitRequest{StreamID: streamID, Payload: []byte("chunk")}) if errMarshal != nil { t.Fatalf("marshal emit request: %v", errMarshal) } if _, errEmit := host.callFromPlugin(context.Background(), pluginabi.MethodHostStreamEmit, emitReq); errEmit != nil { t.Fatalf("emit callback error = %v", errEmit) } closeReq, errMarshal := json.Marshal(rpcStreamCloseRequest{StreamID: streamID}) if errMarshal != nil { t.Fatalf("marshal close request: %v", errMarshal) } if _, errClose := host.callFromPlugin(context.Background(), pluginabi.MethodHostStreamClose, closeReq); errClose != nil { t.Fatalf("close callback error = %v", errClose) } chunk, ok := <-chunks if !ok { t.Fatalf("stream closed before chunk") } if string(chunk.Payload) != "chunk" || chunk.Err != nil { t.Fatalf("chunk = %#v, want payload chunk", chunk) } if _, ok = <-chunks; ok { t.Fatalf("stream remains open after close") } }