ollie/toolsrv/rpc_integration_test.go

933 lines
24 KiB
Go

package toolsrv_test
import (
"context"
"encoding/json"
"io"
"net"
"os"
"path/filepath"
"runtime"
"sync"
"syscall"
"testing"
"time"
"ollie/toolsrv"
)
// --- Test infrastructure ---
// rpcRequest mirrors the JSON-RPC 2.0 request format.
type rpcRequest struct {
JSONRPC string `json:"jsonrpc"`
ID int64 `json:"id"`
Method string `json:"method"`
Params json.RawMessage `json:"params,omitempty"`
}
// rpcResponse mirrors the JSON-RPC 2.0 response format.
type rpcResponse struct {
JSONRPC string `json:"jsonrpc"`
ID int64 `json:"id"`
Result json.RawMessage `json:"result,omitempty"`
Error *rpcError `json:"error,omitempty"`
Stream bool `json:"stream,omitempty"`
}
type rpcError struct {
Code int `json:"code"`
Message string `json:"message"`
}
type outputNotif struct {
Data string `json:"data"`
}
// testServer wraps a toolsrv.Server and provides RPC dispatch for testing.
// This mimics the cmd/ollie-remote serveRPC loop.
type testServer struct {
srv *toolsrv.Server
registry *toolsrv.Registry
cwd string
}
func newTestServer(t *testing.T) *testServer {
t.Helper()
cwd := t.TempDir()
// Create a .git directory to make isGitRepo return true
os.MkdirAll(filepath.Join(cwd, ".git"), 0755)
reg, err := toolsrv.NewRegistry()
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
sessionID := "test-session-" + t.Name()
srv := toolsrv.New(cwd)
srv.SetToolRegistry(reg, sessionID)
return &testServer{
srv: srv,
registry: reg,
cwd: cwd,
}
}
// handleRPC dispatches a single RPC request and returns the response.
// This mirrors cmd/ollie-remote handleRPC but returns instead of encoding.
func (ts *testServer) handleRPC(ctx context.Context, req rpcRequest) rpcResponse {
switch req.Method {
case "ping":
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`"pong"`)}
case "host_info":
info := map[string]any{
"platform": runtime.GOOS,
"arch": runtime.GOARCH,
"is_git_repo": isGitRepo(ts.cwd),
}
result, _ := json.Marshal(info)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "list_tools":
tools, err := ts.srv.ListTools()
if err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}}
}
result, _ := json.Marshal(tools)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "set_env":
var params struct {
Key string `json:"key"`
Value string `json:"value"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
ts.srv.SetEnv(params.Key, params.Value)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)}
case "set_cwd":
var params struct {
Dir string `json:"dir"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
ts.srv.SetCWD(params.Dir)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)}
case "set_allow_tools":
var params struct {
Names []string `json:"names"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
ts.srv.SetAllowTools(params.Names)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)}
case "is_parallel_read":
var params struct {
Name string `json:"name"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
ok := ts.srv.IsParallelRead(params.Name)
result, _ := json.Marshal(ok)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "result_tier":
var params struct {
Name string `json:"name"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
tier := ts.srv.ResultTier(params.Name)
result, _ := json.Marshal(tier)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "result_tier_args":
var params struct {
Name string `json:"name"`
Args json.RawMessage `json:"args"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
tier := ts.srv.ResultTierArgs(params.Name, params.Args)
result, _ := json.Marshal(tier)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "detach":
ok := ts.srv.Detach()
result, _ := json.Marshal(ok)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "list_detached":
list := ts.srv.ListDetachedRaw()
result, _ := json.Marshal(list)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "signal_detached":
var params struct {
PID int `json:"pid"`
Signal int `json:"signal"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
if err := ts.srv.SignalDetached(params.PID, syscall.Signal(params.Signal)); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}}
}
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)}
case "get_detached_output":
var params struct {
PID int `json:"pid"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
out, err := ts.srv.GetDetachedOutput(params.PID)
if err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}}
}
result, _ := json.Marshal(out)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "dismiss_detached":
var params struct {
PID int `json:"pid"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
ok := ts.srv.DismissDetached(params.PID)
result, _ := json.Marshal(ok)
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
case "tool_load":
var params struct {
Name string `json:"name"`
}
if err := json.Unmarshal(req.Params, &params); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}}
}
if err := ts.srv.LoadTool(params.Name); err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}}
}
result, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name})
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
default:
// Unknown method = tool call
result, err := ts.srv.CallTool(ctx, req.Method, req.Params)
if err != nil {
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}}
}
return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result}
}
}
func isGitRepo(dir string) bool {
_, err := os.Stat(filepath.Join(dir, ".git"))
return err == nil
}
// --- Pipe-based RPC client/server ---
// pipeConn creates an in-memory connection for testing RPC client/server interaction.
// Returns client and server ends of a bidirectional pipe.
func pipeConn() (client, server net.Conn) {
clientConn, serverConn := net.Pipe()
return clientConn, serverConn
}
// runTestRPCServer runs the test server's RPC handler over the given connection.
// Blocks until the connection is closed or context is cancelled.
func runTestRPCServer(ctx context.Context, ts *testServer, conn net.Conn) {
dec := json.NewDecoder(conn)
enc := json.NewEncoder(conn)
var encMu sync.Mutex
for {
select {
case <-ctx.Done():
return
default:
}
// Set read deadline to allow periodic context checks
conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
var req rpcRequest
if err := dec.Decode(&req); err != nil {
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
continue // Just a timeout, check context and retry
}
if err == io.EOF || ctx.Err() != nil {
return
}
encMu.Lock()
enc.Encode(rpcResponse{
JSONRPC: "2.0",
ID: req.ID,
Error: &rpcError{Code: -32700, Message: "parse error: " + err.Error()},
})
encMu.Unlock()
continue
}
// Handle request (synchronously for simplicity in tests)
resp := ts.handleRPC(ctx, req)
encMu.Lock()
if err := enc.Encode(resp); err != nil {
encMu.Unlock()
return
}
encMu.Unlock()
}
}
// --- Integration tests ---
func TestRPC_Ping(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "ping"}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result string
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if result != "pong" {
t.Errorf("expected 'pong', got %q", result)
}
}
func TestRPC_HostInfo(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "host_info"}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var info toolsrv.HostInfo
if err := json.Unmarshal(resp.Result, &info); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if info.Platform != runtime.GOOS {
t.Errorf("platform: expected %q, got %q", runtime.GOOS, info.Platform)
}
if info.Arch != runtime.GOARCH {
t.Errorf("arch: expected %q, got %q", runtime.GOARCH, info.Arch)
}
if !info.IsGitRepo {
t.Error("expected is_git_repo=true (we created .git dir)")
}
}
func TestRPC_ListTools_Empty(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "list_tools"}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var tools []toolsrv.ToolInfo
if err := json.Unmarshal(resp.Result, &tools); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
// No tools loaded yet
if len(tools) != 0 {
t.Errorf("expected 0 tools, got %d", len(tools))
}
}
func TestRPC_SetEnv(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
// Hook to verify env is set
var envKey, envValue string
ts.srv.OnEnvSet = func(k, v string) {
envKey, envValue = k, v
}
ctx := context.Background()
params, _ := json.Marshal(map[string]string{"key": "TEST_VAR", "value": "test_value"})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_env", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result bool
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if !result {
t.Error("expected result=true")
}
if envKey != "TEST_VAR" || envValue != "test_value" {
t.Errorf("env not set: got key=%q value=%q", envKey, envValue)
}
}
func TestRPC_SetCWD(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
newDir := "/tmp/test-cwd"
ctx := context.Background()
params, _ := json.Marshal(map[string]string{"dir": newDir})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_cwd", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result bool
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if !result {
t.Error("expected result=true")
}
}
func TestRPC_SetAllowTools(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]any{"names": []string{"tool1", "tool2"}})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_allow_tools", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result bool
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if !result {
t.Error("expected result=true")
}
// Verify allowed tools are set
allowed := ts.srv.AllowTools()
if len(allowed) != 2 {
t.Errorf("expected 2 allowed tools, got %d", len(allowed))
}
}
func TestRPC_IsParallelRead(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]string{"name": "unknown_tool"})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "is_parallel_read", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result bool
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
// Unknown tools default to not parallel-read-safe
if result {
t.Error("expected false for unknown tool")
}
}
func TestRPC_ResultTier(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]string{"name": "unknown_tool"})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "result_tier", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var tier string
if err := json.Unmarshal(resp.Result, &tier); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
// Unknown tools default to "hot"
if tier != "hot" {
t.Errorf("expected tier='hot', got %q", tier)
}
}
func TestRPC_ResultTierArgs(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]any{
"name": "some_tool",
"args": json.RawMessage(`{"foo": "bar"}`),
})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "result_tier_args", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var tier string
if err := json.Unmarshal(resp.Result, &tier); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
// Unknown tools default to "hot"
if tier != "hot" {
t.Errorf("expected tier='hot', got %q", tier)
}
}
func TestRPC_Detach_NoRunningProcess(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "detach"}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result bool
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
// No running process to detach
if result {
t.Error("expected false when no process running")
}
}
func TestRPC_ListDetached_Empty(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "list_detached"}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var list []any
if err := json.Unmarshal(resp.Result, &list); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
if len(list) != 0 {
t.Errorf("expected empty list, got %d items", len(list))
}
}
func TestRPC_SignalDetached_NotFound(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]any{"pid": 99999, "signal": int(syscall.SIGTERM)})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "signal_detached", Params: params}
resp := ts.handleRPC(ctx, req)
// Should error because no such detached process
if resp.Error == nil {
t.Fatal("expected error for nonexistent PID")
}
if resp.Error.Code != -32000 {
t.Errorf("expected error code -32000, got %d", resp.Error.Code)
}
}
func TestRPC_GetDetachedOutput_NotFound(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]any{"pid": 99999})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "get_detached_output", Params: params}
resp := ts.handleRPC(ctx, req)
// Should error because no such detached process
if resp.Error == nil {
t.Fatal("expected error for nonexistent PID")
}
}
func TestRPC_DismissDetached_NotFound(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]any{"pid": 99999})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "dismiss_detached", Params: params}
resp := ts.handleRPC(ctx, req)
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
var result bool
if err := json.Unmarshal(resp.Result, &result); err != nil {
t.Fatalf("unmarshal result: %v", err)
}
// Should return false for nonexistent PID
if result {
t.Error("expected false for nonexistent PID")
}
}
func TestRPC_CallTool_NotLoaded(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
params, _ := json.Marshal(map[string]any{"arg": "value"})
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "nonexistent_tool", Params: params}
resp := ts.handleRPC(ctx, req)
// Should error because tool not loaded
if resp.Error == nil {
t.Fatal("expected error for unloaded tool")
}
if resp.Error.Code != -32000 {
t.Errorf("expected error code -32000, got %d", resp.Error.Code)
}
}
func TestRPC_InvalidParams(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx := context.Background()
// Send malformed JSON as params
req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_env", Params: json.RawMessage(`{invalid`)}
resp := ts.handleRPC(ctx, req)
// Should error with invalid params
if resp.Error == nil {
t.Fatal("expected error for invalid params")
}
if resp.Error.Code != -32602 {
t.Errorf("expected error code -32602 (invalid params), got %d", resp.Error.Code)
}
}
// --- Full pipe-based client/server tests ---
func TestRPC_PipeConnection(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
clientConn, serverConn := pipeConn()
defer clientConn.Close()
defer serverConn.Close()
// Start server in background
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
runTestRPCServer(ctx, ts, serverConn)
}()
// Create client connection using toolsrv.Conn
conn := toolsrv.NewConn(clientConn, nil)
// Test ping
if err := conn.Ping(); err != nil {
t.Errorf("Ping failed: %v", err)
}
// Test list tools (should be empty)
tools, err := conn.ListTools()
if err != nil {
t.Errorf("ListTools failed: %v", err)
}
if len(tools) != 0 {
t.Errorf("expected 0 tools, got %d", len(tools))
}
// Test set env
conn.SetEnv("TEST_KEY", "TEST_VALUE")
// Test set cwd
conn.SetCWD("/tmp/test")
// Test is_parallel_read
if conn.IsParallelRead("unknown") {
t.Error("expected IsParallelRead=false for unknown tool")
}
// Test result_tier
tier := conn.ResultTier("unknown")
if tier != "hot" {
t.Errorf("expected tier='hot', got %q", tier)
}
// Clean up
cancel()
conn.Close()
wg.Wait()
}
func TestRPC_ConcurrentRequests(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
clientConn, serverConn := pipeConn()
defer clientConn.Close()
defer serverConn.Close()
// Start server
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
runTestRPCServer(ctx, ts, serverConn)
}()
conn := toolsrv.NewConn(clientConn, nil)
// Fire off concurrent ping requests
const numRequests = 10
var reqWg sync.WaitGroup
errors := make(chan error, numRequests)
for i := 0; i < numRequests; i++ {
reqWg.Add(1)
go func() {
defer reqWg.Done()
if err := conn.Ping(); err != nil {
errors <- err
}
}()
}
reqWg.Wait()
close(errors)
for err := range errors {
t.Errorf("concurrent ping failed: %v", err)
}
cancel()
conn.Close()
wg.Wait()
}
func TestRPC_ContextCancellation(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx, cancel := context.WithCancel(context.Background())
clientConn, serverConn := pipeConn()
defer clientConn.Close()
defer serverConn.Close()
// Start server
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
runTestRPCServer(ctx, ts, serverConn)
}()
conn := toolsrv.NewConn(clientConn, nil)
// Cancel context
cancel()
// Allow some time for cancellation to propagate
time.Sleep(50 * time.Millisecond)
// Clean up
conn.Close()
wg.Wait()
}
func TestRPC_ConnectionClose(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
clientConn, serverConn := pipeConn()
// Start server
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
runTestRPCServer(ctx, ts, serverConn)
}()
conn := toolsrv.NewConn(clientConn, nil)
// Verify ping works
if err := conn.Ping(); err != nil {
t.Fatalf("initial Ping failed: %v", err)
}
// Close the client connection - this should cause the server to detect closure
conn.Close()
clientConn.Close()
serverConn.Close()
// Cancel context to ensure server exits
cancel()
// Wait with timeout
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
// Good, server exited
case <-time.After(2 * time.Second):
t.Error("server did not exit after connection close")
}
}
// TestRPC_JSONRPCProtocol tests raw JSON-RPC protocol compliance
func TestRPC_JSONRPCProtocol(t *testing.T) {
ts := newTestServer(t)
defer ts.srv.Close()
tests := []struct {
name string
request rpcRequest
wantErr bool
errCode int
checkResult func(t *testing.T, result json.RawMessage)
}{
{
name: "ping returns pong",
request: rpcRequest{JSONRPC: "2.0", ID: 1, Method: "ping"},
checkResult: func(t *testing.T, result json.RawMessage) {
var s string
json.Unmarshal(result, &s)
if s != "pong" {
t.Errorf("expected 'pong', got %q", s)
}
},
},
{
name: "set_env with valid params",
request: rpcRequest{
JSONRPC: "2.0",
ID: 2,
Method: "set_env",
Params: json.RawMessage(`{"key":"K","value":"V"}`),
},
checkResult: func(t *testing.T, result json.RawMessage) {
var b bool
json.Unmarshal(result, &b)
if !b {
t.Error("expected true")
}
},
},
{
name: "set_env with invalid params",
request: rpcRequest{
JSONRPC: "2.0",
ID: 3,
Method: "set_env",
Params: json.RawMessage(`not valid json`),
},
wantErr: true,
errCode: -32602,
},
{
name: "unknown tool returns error",
request: rpcRequest{
JSONRPC: "2.0",
ID: 4,
Method: "fake_tool",
Params: json.RawMessage(`{}`),
},
wantErr: true,
errCode: -32000,
},
}
ctx := context.Background()
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resp := ts.handleRPC(ctx, tt.request)
if tt.wantErr {
if resp.Error == nil {
t.Fatal("expected error, got none")
}
if resp.Error.Code != tt.errCode {
t.Errorf("error code: expected %d, got %d", tt.errCode, resp.Error.Code)
}
return
}
if resp.Error != nil {
t.Fatalf("unexpected error: %v", resp.Error.Message)
}
if tt.checkResult != nil {
tt.checkResult(t, resp.Result)
}
})
}
}