137 lines
3.1 KiB
Go
137 lines
3.1 KiB
Go
package rhntest
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"testing"
|
|
)
|
|
|
|
// Server is an httptest Robinhood MCP mock that speaks JSON-RPC tools/call.
|
|
type Server struct {
|
|
URL string
|
|
Token string // if non-empty, require Authorization: Bearer <Token>
|
|
|
|
mu sync.Mutex
|
|
srv *httptest.Server
|
|
results map[string]json.RawMessage
|
|
rpcErrs map[string]string
|
|
httpStatus int
|
|
httpBody string
|
|
lastName string
|
|
lastArgs map[string]any
|
|
}
|
|
|
|
// New starts an httptest server and registers t.Cleanup(Close).
|
|
func New(t *testing.T) *Server {
|
|
t.Helper()
|
|
s := &Server{
|
|
results: make(map[string]json.RawMessage),
|
|
rpcErrs: make(map[string]string),
|
|
}
|
|
s.srv = httptest.NewServer(http.HandlerFunc(s.serve))
|
|
s.URL = s.srv.URL
|
|
t.Cleanup(s.Close)
|
|
return s
|
|
}
|
|
|
|
// Set registers a JSON-RPC result for a tool name.
|
|
func (s *Server) Set(name string, result json.RawMessage) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.results[name] = result
|
|
}
|
|
|
|
// SetHTTPError makes every request return this HTTP status and body.
|
|
func (s *Server) SetHTTPError(status int, body string) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.httpStatus = status
|
|
s.httpBody = body
|
|
}
|
|
|
|
// SetRPCError makes tools/call for name return a JSON-RPC error.
|
|
func (s *Server) SetRPCError(name, message string) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.rpcErrs[name] = message
|
|
}
|
|
|
|
// LastName is the most recent tools/call params.name.
|
|
func (s *Server) LastName() string {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return s.lastName
|
|
}
|
|
|
|
// LastArgs is the most recent tools/call params.arguments.
|
|
func (s *Server) LastArgs() map[string]any {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return s.lastArgs
|
|
}
|
|
|
|
// Close shuts down the httptest server.
|
|
func (s *Server) Close() {
|
|
s.srv.Close()
|
|
}
|
|
|
|
type rpcError struct {
|
|
Message string `json:"message"`
|
|
}
|
|
|
|
type rpcReply struct {
|
|
JSONRPC string `json:"jsonrpc"`
|
|
ID json.RawMessage `json:"id"`
|
|
Result json.RawMessage `json:"result,omitempty"`
|
|
Error *rpcError `json:"error,omitempty"`
|
|
}
|
|
|
|
func (s *Server) serve(w http.ResponseWriter, r *http.Request) {
|
|
if s.Token != "" && r.Header.Get("Authorization") != "Bearer "+s.Token {
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
s.mu.Lock()
|
|
httpStatus := s.httpStatus
|
|
httpBody := s.httpBody
|
|
s.mu.Unlock()
|
|
if httpStatus != 0 {
|
|
w.WriteHeader(httpStatus)
|
|
_, _ = w.Write([]byte(httpBody))
|
|
return
|
|
}
|
|
|
|
var req struct {
|
|
ID json.RawMessage `json:"id"`
|
|
Params struct {
|
|
Name string `json:"name"`
|
|
Arguments map[string]any `json:"arguments"`
|
|
} `json:"params"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
s.mu.Lock()
|
|
s.lastName = req.Params.Name
|
|
s.lastArgs = req.Params.Arguments
|
|
rpcMsg, hasRPC := s.rpcErrs[req.Params.Name]
|
|
result, hasResult := s.results[req.Params.Name]
|
|
s.mu.Unlock()
|
|
|
|
reply := rpcReply{JSONRPC: "2.0", ID: req.ID}
|
|
if hasRPC {
|
|
reply.Error = &rpcError{Message: rpcMsg}
|
|
} else if !hasResult {
|
|
reply.Error = &rpcError{Message: "unknown tool"}
|
|
} else {
|
|
reply.Result = result
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(reply)
|
|
}
|