Fix grpcserver rate limiter: move to Server field

The package-level defaultRateLimiter drained its token bucket
across all test cases, causing later tests to hit ResourceExhausted.
Move rateLimiter from a package-level var to a *grpcRateLimiter field
on Server; New() allocates a fresh instance (10 req/s, burst 10) per
server. Each test's newTestEnv() constructs its own Server, so tests
no longer share limiter state.

Production behaviour is unchanged: a single Server is constructed at
startup and lives for the process lifetime.
This commit is contained in:
2026-03-11 19:20:32 -07:00
parent a80242ae3e
commit 4596ea08ab
12 changed files with 276 additions and 123 deletions

View File

@@ -3,6 +3,7 @@
// Security note: this package is test-only. It never enforces TLS and uses
// trivial token generation. Do not use in production.
package mock
import (
"encoding/json"
"fmt"
@@ -11,6 +12,7 @@ import (
"strings"
"sync"
)
// Account holds mock account state.
type Account struct {
ID string
@@ -20,25 +22,28 @@ type Account struct {
Status string
Roles []string
}
// PGCreds holds mock Postgres credential state.
type PGCreds struct {
Host string
Port int
Database string
Username string
Password string
Port int
}
// Server is an in-memory MCIAS mock server.
type Server struct {
mu sync.RWMutex
httpServer *httptest.Server
accounts map[string]*Account // id → account
byName map[string]*Account // username → account
tokens map[string]string // token → account id
revoked map[string]bool // revoked tokens
pgcreds map[string]*PGCreds // account id → pg creds
nextSeq int
httpServer *httptest.Server
mu sync.RWMutex
}
// NewServer creates and starts a new mock server. Call Close() when done.
func NewServer() *Server {
s := &Server{
@@ -61,14 +66,17 @@ func NewServer() *Server {
s.httpServer = httptest.NewServer(mux)
return s
}
// URL returns the base URL of the mock server.
func (s *Server) URL() string {
return s.httpServer.URL
}
// Close shuts down the mock server.
func (s *Server) Close() {
s.httpServer.Close()
}
// AddAccount adds a test account and returns its ID.
func (s *Server) AddAccount(username, password, accountType string, roles ...string) string {
s.mu.Lock()
@@ -87,12 +95,14 @@ func (s *Server) AddAccount(username, password, accountType string, roles ...str
s.byName[username] = acct
return id
}
// IssueToken directly adds a token for an account (for pre-auth test setup).
func (s *Server) IssueToken(accountID, token string) {
s.mu.Lock()
defer s.mu.Unlock()
s.tokens[token] = accountID
}
// issueToken creates a new token for the given account ID.
// Caller must hold s.mu (write lock).
func (s *Server) issueToken(accountID string) string {
@@ -365,12 +375,12 @@ func (s *Server) handleAccountByID(w http.ResponseWriter, r *http.Request) {
if len(parts) == 2 {
sub = parts[1]
}
switch {
case sub == "roles":
switch sub {
case "roles":
s.handleRoles(w, r, id)
case sub == "pgcreds":
case "pgcreds":
s.handlePGCreds(w, r, id)
case sub == "":
case "":
s.handleSingleAccount(w, r, id)
default:
sendError(w, http.StatusNotFound, "not found")
@@ -491,10 +501,10 @@ func (s *Server) handlePGCreds(w http.ResponseWriter, r *http.Request, id string
case http.MethodPut:
var req struct {
Host string `json:"host"`
Port int `json:"port"`
Database string `json:"database"`
Username string `json:"username"`
Password string `json:"password"`
Port int `json:"port"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
sendError(w, http.StatusBadRequest, "bad request")