Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
2.2 kB · 76 lines
Go
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677package pluginhost
import ( "context" "encoding/json" "net/http" "testing" "time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi")
func TestHostModelExecuteStreamDetachesFromCallbackParentCancel(t *testing.T) { host := New() ctxSeen := make(chan context.Context, 1) host.SetModelExecutor(&fakeHostModelExecutor{ executeModelStream: func(ctx context.Context, req handlers.ModelExecutionRequest) (handlers.ModelExecutionStream, *interfaces.ErrorMessage) { ctxSeen <- ctx return handlers.ModelExecutionStream{ StatusCode: http.StatusOK, Chunks: make(chan handlers.ModelExecutionChunk), }, nil }, }) parentCtx, cancelParent := context.WithCancel(context.Background()) callbackID, closeCallback := host.openCallbackContext(parentCtx) defer closeCallback()
rawReq, errMarshal := json.Marshal(rpcHostModelExecutionRequest{ HostModelExecutionRequest: pluginapi.HostModelExecutionRequest{ EntryProtocol: "openai", ExitProtocol: "openai", Model: "model-1", Stream: true, Body: []byte(`{"stream":true}`), }, HostCallbackID: callbackID, }) if errMarshal != nil { t.Fatalf("marshal request: %v", errMarshal) } rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostModelExecuteStream, rawReq) if errCall != nil { t.Fatalf("callFromPlugin() error = %v", errCall) } resp, errDecode := decodeRPCEnvelope[pluginapi.HostModelStreamResponse](rawResp) if errDecode != nil { t.Fatalf("decode response: %v", errDecode) } if resp.StreamID == "" { t.Fatalf("stream id is empty: %#v", resp) }
var streamCtx context.Context select { case streamCtx = <-ctxSeen: case <-time.After(time.Second): t.Fatal("model executor was not called") } cancelParent() select { case <-streamCtx.Done(): t.Fatal("stream context was canceled by callback parent context") default: }
closeCallback() select { case <-streamCtx.Done(): case <-time.After(time.Second): t.Fatal("stream context was not canceled after callback scope closed") }}