164 lines
5.6 KiB
Go
164 lines
5.6 KiB
Go
package pluginhost
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
|
)
|
|
|
|
func TestRequestInterceptorTerminationStopsChain(t *testing.T) {
|
|
lowCalls := 0
|
|
host := newHostWithRecords(
|
|
capabilityRecord{
|
|
id: "high",
|
|
priority: 20,
|
|
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{
|
|
RequestInterceptor: requestInterceptorFunc(func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) {
|
|
return pluginapi.RequestInterceptResponse{
|
|
Terminate: true,
|
|
StatusCode: http.StatusForbidden,
|
|
ResponseHeaders: http.Header{"Content-Type": {"application/json"}},
|
|
ResponseBody: []byte(`{"error":"blocked"}`),
|
|
}, nil
|
|
}),
|
|
}},
|
|
},
|
|
capabilityRecord{
|
|
id: "low",
|
|
priority: 10,
|
|
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{
|
|
RequestInterceptor: requestInterceptorFunc(func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) {
|
|
lowCalls++
|
|
return pluginapi.RequestInterceptResponse{}, nil
|
|
}),
|
|
}},
|
|
},
|
|
)
|
|
|
|
response := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{RequestID: "request-1"})
|
|
if !response.Terminate || response.StatusCode != http.StatusForbidden {
|
|
t.Fatalf("termination response = %#v", response)
|
|
}
|
|
if response.ResponseHeaders.Get("Content-Type") != "application/json" || string(response.ResponseBody) != `{"error":"blocked"}` {
|
|
t.Fatalf("termination payload = %#v", response)
|
|
}
|
|
if lowCalls != 0 {
|
|
t.Fatalf("lower-priority interceptor calls = %d, want 0", lowCalls)
|
|
}
|
|
}
|
|
|
|
func TestCompleteRequestUsesUncancelledContextAndClonesMetadata(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
originalNested := map[string]any{"value": "original"}
|
|
var got pluginapi.RequestCompletion
|
|
var callbackContextError error
|
|
done := make(chan struct{})
|
|
host := newHostWithRecords(capabilityRecord{
|
|
id: "lifecycle",
|
|
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{
|
|
RequestLifecyclePlugin: requestLifecyclePluginFunc(func(callbackCtx context.Context, completion pluginapi.RequestCompletion) {
|
|
callbackContextError = callbackCtx.Err()
|
|
got = completion
|
|
completion.Metadata["nested"].(map[string]any)["value"] = "mutated"
|
|
close(done)
|
|
}),
|
|
}},
|
|
})
|
|
|
|
host.CompleteRequest(ctx, pluginapi.RequestCompletion{
|
|
RequestID: "request-1",
|
|
Outcome: pluginapi.RequestCompletionCanceled,
|
|
StartedAt: time.Now().Add(-time.Second),
|
|
CompletedAt: time.Now(),
|
|
Metadata: map[string]any{"nested": originalNested},
|
|
})
|
|
<-done
|
|
|
|
if callbackContextError != nil {
|
|
t.Fatalf("callback context error = %v", callbackContextError)
|
|
}
|
|
if got.RequestID != "request-1" || got.Outcome != pluginapi.RequestCompletionCanceled {
|
|
t.Fatalf("completion = %#v", got)
|
|
}
|
|
if originalNested["value"] != "original" {
|
|
t.Fatalf("input metadata was mutated: %#v", originalNested)
|
|
}
|
|
}
|
|
|
|
func TestCompleteRequestDoesNotWaitForBlockingPlugin(t *testing.T) {
|
|
started := make(chan struct{})
|
|
release := make(chan struct{})
|
|
host := newHostWithRecords(capabilityRecord{
|
|
id: "blocking-lifecycle",
|
|
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{
|
|
RequestLifecyclePlugin: requestLifecyclePluginFunc(func(context.Context, pluginapi.RequestCompletion) {
|
|
close(started)
|
|
<-release
|
|
}),
|
|
}},
|
|
})
|
|
|
|
returned := make(chan struct{})
|
|
go func() {
|
|
host.CompleteRequest(context.Background(), pluginapi.RequestCompletion{RequestID: "request-blocking"})
|
|
close(returned)
|
|
}()
|
|
select {
|
|
case <-returned:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("CompleteRequest blocked on lifecycle plugin")
|
|
}
|
|
select {
|
|
case <-started:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("lifecycle plugin was not invoked")
|
|
}
|
|
close(release)
|
|
}
|
|
|
|
func TestRPCCapabilitiesAndAdapterIncludeRequestLifecycle(t *testing.T) {
|
|
var got pluginapi.RequestCompletion
|
|
plugin := validTestPlugin("request-lifecycle")
|
|
plugin.Capabilities.RequestLifecyclePlugin = requestLifecyclePluginFunc(func(_ context.Context, completion pluginapi.RequestCompletion) {
|
|
got = completion
|
|
})
|
|
caps := rpcCapabilitiesFromPlugin(plugin)
|
|
if !caps.RequestLifecyclePlugin {
|
|
t.Fatal("RequestLifecyclePlugin = false, want true")
|
|
}
|
|
rawCaps, errMarshal := json.Marshal(caps)
|
|
if errMarshal != nil {
|
|
t.Fatalf("Marshal() error = %v", errMarshal)
|
|
}
|
|
var decoded map[string]any
|
|
if errUnmarshal := json.Unmarshal(rawCaps, &decoded); errUnmarshal != nil {
|
|
t.Fatalf("Unmarshal() error = %v", errUnmarshal)
|
|
}
|
|
if decoded["request_lifecycle_plugin"] != true {
|
|
t.Fatalf("request_lifecycle_plugin = %#v", decoded["request_lifecycle_plugin"])
|
|
}
|
|
|
|
lookup := newTestSymbolLookup(&testPlugin{registerResult: plugin})
|
|
registered, errRegister := registerRPCPlugin(context.Background(), nil, "request-lifecycle", lookup, pluginabi.MethodPluginRegister, nil)
|
|
if errRegister != nil {
|
|
t.Fatalf("registerRPCPlugin() error = %v", errRegister)
|
|
}
|
|
if registered.Capabilities.RequestLifecyclePlugin == nil {
|
|
t.Fatal("RequestLifecyclePlugin = nil, want RPC adapter")
|
|
}
|
|
if errComplete := registered.Capabilities.RequestLifecyclePlugin.HandleRequestComplete(context.Background(), pluginapi.RequestCompletion{
|
|
RequestID: "request-rpc",
|
|
Outcome: pluginapi.RequestCompletionSucceeded,
|
|
}); errComplete != nil {
|
|
t.Fatalf("HandleRequestComplete() error = %v", errComplete)
|
|
}
|
|
if got.RequestID != "request-rpc" || got.Outcome != pluginapi.RequestCompletionSucceeded {
|
|
t.Fatalf("RPC completion = %#v", got)
|
|
}
|
|
}
|