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:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user