Alassane
cc33b2a813
Co-authored-by: ayattara <alassane.yattara@savoirfairelinux.com> Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
77 lines
2.5 KiB
Go
77 lines
2.5 KiB
Go
package pluginsdk
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"testing"
|
|
|
|
pluginv1 "github.com/kandev/kandev/proto/kandev/plugin/v1"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/grpc/test/bufconn"
|
|
)
|
|
|
|
type utilityRPCProbeServer struct {
|
|
pluginv1.UnimplementedHostServer
|
|
oldCalls int
|
|
optionsCalls int
|
|
}
|
|
|
|
func (s *utilityRPCProbeServer) InvokeUtilityAgent(context.Context, *pluginv1.InvokeUtilityAgentRequest) (*pluginv1.InvokeUtilityAgentResponse, error) {
|
|
s.oldCalls++
|
|
return &pluginv1.InvokeUtilityAgentResponse{Text: "old/default"}, nil
|
|
}
|
|
|
|
func (s *utilityRPCProbeServer) InvokeUtilityAgentWithOptions(context.Context, *pluginv1.InvokeUtilityAgentWithOptionsRequest) (*pluginv1.InvokeUtilityAgentResponse, error) {
|
|
s.optionsCalls++
|
|
return nil, status.Error(codes.Unimplemented, "new rpc unavailable")
|
|
}
|
|
|
|
func dialRawHostClient(t *testing.T, register func(*grpc.Server)) (pluginv1.HostClient, Host) {
|
|
t.Helper()
|
|
|
|
lis := bufconn.Listen(1024 * 1024)
|
|
srv := grpc.NewServer()
|
|
register(srv)
|
|
go func() { _ = srv.Serve(lis) }()
|
|
t.Cleanup(srv.Stop)
|
|
|
|
conn, err := grpc.NewClient("passthrough:///bufconn",
|
|
grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
|
|
return lis.DialContext(ctx)
|
|
}),
|
|
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
|
)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { _ = conn.Close() })
|
|
|
|
return pluginv1.NewHostClient(conn), newHostClient(conn)
|
|
}
|
|
|
|
func TestInvokeUtilityAgent_ExplicitOverrideDoesNotDowngradeOnUnimplemented(t *testing.T) {
|
|
probe := &utilityRPCProbeServer{}
|
|
_, host := dialRawHostClient(t, func(server *grpc.Server) {
|
|
pluginv1.RegisterHostServer(server, probe)
|
|
})
|
|
|
|
_, err := host.InvokeUtilityAgent(context.Background(), "prompt", UtilityAgentOptions{ProfileID: "profile-1"})
|
|
require.Equal(t, codes.Unimplemented, status.Code(err))
|
|
require.Equal(t, 1, probe.optionsCalls)
|
|
require.Zero(t, probe.oldCalls)
|
|
}
|
|
|
|
func TestInvokeUtilityAgent_LegacyWireUsesDefault(t *testing.T) {
|
|
impl := &dataRecordingHost{utilityText: "default completion"}
|
|
client, _ := dialRawHostClient(t, func(server *grpc.Server) {
|
|
registerHostServer(server, impl)
|
|
})
|
|
|
|
response, err := client.InvokeUtilityAgent(context.Background(), &pluginv1.InvokeUtilityAgentRequest{Prompt: "prompt"})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "default completion", response.GetText())
|
|
require.Empty(t, impl.lastUtilityOptions)
|
|
}
|