Janady13/missionmatch / internal/platform/federated/federated_test.go

1package federated 2 3// federated_test.go exercises the HARD spend gate (Phase 3) with in memory fakes for the ledger store and the 4// audit seam, so the gate behavior is proven with NO database connected: an in budget release persists and is 5// audited, an over budget release is refused and writes nothing, the accumulation is read from the persisted 6// ledger never from the caller (never client trusted), and the Clip step is deterministic. The flag off 7// refusal is covered in accountant_test.go. 8 9import ( 10 "context" 11 "encoding/json" 12 "errors" 13 "testing" 14 15 "github.com/google/uuid" 16 17 "github.com/missionworkforce/missionmatch/internal/platform/audit" 18 "github.com/missionworkforce/missionmatch/internal/platform/config" 19 "github.com/missionworkforce/missionmatch/internal/platform/database/dbgen" 20) 21 22// fakeLedger is an in memory LedgerStore. It records inserts and returns the latest as the running budget, so 23// the gate composes across releases through the persisted (not caller supplied) accumulation. 24type fakeLedger struct { 25 rows []dbgen.DpBudgetLedger 26 insertErr error // when set, Insert fails (to prove a failed write is not audited) 27 latestErr error // when set, LatestForKey fails 28 insertSeen []dbgen.InsertDpBudgetLedgerParams 29} 30 31func (f *fakeLedger) LatestForKey(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (dbgen.DpBudgetLedger, bool, error) { 32 if f.latestErr != nil { 33 return dbgen.DpBudgetLedger{}, false, f.latestErr 34 } 35 // Return the most recent row that matches the budget key, mirroring the SQL coalesce key match. 36 var match *dbgen.DpBudgetLedger 37 for i := range f.rows { 38 r := f.rows[i] 39 if r.ModelVersion != modelVersion { 40 continue 41 } 42 if sameTenantKey(r.TenantID, tenantID) { 43 match = &f.rows[i] 44 } 45 } 46 if match == nil { 47 return dbgen.DpBudgetLedger{}, false, nil 48 } 49 return *match, true, nil 50} 51 52func (f *fakeLedger) Insert(ctx context.Context, arg dbgen.InsertDpBudgetLedgerParams) (dbgen.DpBudgetLedger, error) { 53 f.insertSeen = append(f.insertSeen, arg) 54 if f.insertErr != nil { 55 return dbgen.DpBudgetLedger{}, f.insertErr 56 } 57 row := dbgen.DpBudgetLedger{ 58 ID: uuid.New(), 59 ModelVersion: arg.ModelVersion, 60 TenantID: arg.TenantID, 61 RdpOrders: arg.RdpOrders, 62 AccumulatedRdp: arg.AccumulatedRdp, 63 EpsilonAtDelta: arg.EpsilonAtDelta, 64 Delta: arg.Delta, 65 ContributingTenantCount: arg.ContributingTenantCount, 66 } 67 f.rows = append(f.rows, row) 68 return row, nil 69} 70 71func sameTenantKey(a, b *uuid.UUID) bool { 72 if a == nil && b == nil { 73 return true 74 } 75 if a == nil || b == nil { 76 return false 77 } 78 return *a == *b 79} 80 81// fakeAuditor records the events the gate emits so the test asserts a release is audited as 82// federated.dp.released with the right payload and no PII. 83type fakeAuditor struct { 84 events []audit.Event 85 failErr error 86} 87 88func (f *fakeAuditor) System(ctx context.Context, e audit.Event) error { 89 if f.failErr != nil { 90 return f.failErr 91 } 92 f.events = append(f.events, e) 93 return nil 94} 95 96// newTestGate builds a gate over the fakes, flag ON, with a generous epsilon so the in budget path runs. 97func newTestGate(store LedgerStore, auditW systemAuditor, epsilon, delta float64) *SpendGate { 98 return &SpendGate{store: store, auditW: auditW, epsilon: epsilon, delta: delta, enabled: true} 99} 100 101// TestGateInBudgetReleasePersistsAndAudits proves an in budget release composes, persists exactly one ledger 102// row through the store, and is audited as federated.dp.released with the model version, the contributing 103// tenant count, and the converted epsilon and delta, and NO PII. 104func TestGateInBudgetReleasePersistsAndAudits(t *testing.T) { 105 store := &fakeLedger{} 106 au := &fakeAuditor{} 107 gate := newTestGate(store, au, 100, 1e-6) 108 109 rel := Release{ 110 ModelVersion: "v1.4.0", 111 TenantID: nil, // cross tenant cohort release 112 ContributingTenantCount: 7, 113 Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}, 114 } 115 eps, err := gate.Charge(context.Background(), rel) 116 if err != nil { 117 t.Fatalf("in budget Charge failed: %v", err) 118 } 119 if eps <= 0 { 120 t.Fatalf("expected a positive epsilon, got %v", eps) 121 } 122 if len(store.rows) != 1 { 123 t.Fatalf("expected exactly one persisted ledger row, got %d", len(store.rows)) 124 } 125 if got := store.rows[0].ContributingTenantCount; got != 7 { 126 t.Fatalf("contributing tenant count not recorded: got %d, want 7", got) 127 } 128 if len(au.events) != 1 { 129 t.Fatalf("expected exactly one audit event, got %d", len(au.events)) 130 } 131 ev := au.events[0] 132 if ev.Action != "federated.dp.released" { 133 t.Fatalf("audit action = %q, want federated.dp.released", ev.Action) 134 } 135 var payload map[string]any 136 if err := json.Unmarshal(ev.NewValues, &payload); err != nil { 137 t.Fatalf("audit payload not json: %v", err) 138 } 139 if payload["model_version"] != "v1.4.0" { 140 t.Fatalf("audit payload model_version = %v, want v1.4.0", payload["model_version"]) 141 } 142 // No PII, no candidate id, no raw vector in the payload. 143 for _, banned := range []string{"candidate_id", "name", "email", "resume", "vector"} { 144 if _, ok := payload[banned]; ok { 145 t.Fatalf("audit payload leaked a banned field: %q", banned) 146 } 147 } 148} 149 150// TestGateComposesAcrossReleasesFromLedger proves the gate reads the PERSISTED accumulation (never a caller 151// supplied one) and composes across releases: the epsilon after the second release is strictly greater than 152// after the first, because the accumulation grew in the ledger. 153func TestGateComposesAcrossReleasesFromLedger(t *testing.T) { 154 store := &fakeLedger{} 155 au := &fakeAuditor{} 156 gate := newTestGate(store, au, 1e9, 1e-6) // huge budget so neither release is refused 157 158 rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}} 159 eps1, err := gate.Charge(context.Background(), rel) 160 if err != nil { 161 t.Fatalf("first Charge: %v", err) 162 } 163 eps2, err := gate.Charge(context.Background(), rel) 164 if err != nil { 165 t.Fatalf("second Charge: %v", err) 166 } 167 if eps2 <= eps1 { 168 t.Fatalf("epsilon did not grow across composed releases read from the ledger: %v then %v", eps1, eps2) 169 } 170 if len(store.rows) != 2 { 171 t.Fatalf("expected two persisted rows, got %d", len(store.rows)) 172 } 173} 174 175// TestGateOverBudgetRefusesAndWritesNothing proves an over budget release is refused with ErrBudgetExceeded, 176// persists NO ledger row, and emits NO audit event. The budget is tiny so a single noisy release blows it. 177func TestGateOverBudgetRefusesAndWritesNothing(t *testing.T) { 178 store := &fakeLedger{} 179 au := &fakeAuditor{} 180 gate := newTestGate(store, au, 0.5, 1e-7) // tiny budget 181 182 rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 0.1, ClipNorm: 1}} // huge RDP 183 _, err := gate.Charge(context.Background(), rel) 184 if !errors.Is(err, ErrBudgetExceeded) { 185 t.Fatalf("expected ErrBudgetExceeded, got %v", err) 186 } 187 if len(store.rows) != 0 || len(store.insertSeen) != 0 { 188 t.Fatalf("an over budget release wrote to the ledger: rows=%d inserts=%d", len(store.rows), len(store.insertSeen)) 189 } 190 if len(au.events) != 0 { 191 t.Fatalf("an over budget release was audited: %d events", len(au.events)) 192 } 193} 194 195// TestGateBudgetExhaustsAfterRepeatedReleases proves the budget is genuinely spent down: repeated in budget 196// releases eventually trip the hard gate, after which every further release is refused and writes nothing. 197func TestGateBudgetExhaustsAfterRepeatedReleases(t *testing.T) { 198 store := &fakeLedger{} 199 au := &fakeAuditor{} 200 gate := newTestGate(store, au, 8, 1e-7) // moderate budget, reference Gboard band 201 202 rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 1.5, ClipNorm: 1}} 203 released, refused := 0, 0 204 for i := 0; i < 1000; i++ { 205 _, err := gate.Charge(context.Background(), rel) 206 switch { 207 case err == nil: 208 released++ 209 case errors.Is(err, ErrBudgetExceeded): 210 refused++ 211 default: 212 t.Fatalf("unexpected error at iteration %d: %v", i, err) 213 } 214 } 215 if released == 0 { 216 t.Fatalf("budget never admitted a single release") 217 } 218 if refused == 0 { 219 t.Fatalf("budget never exhausted, the hard gate never fired") 220 } 221 if len(store.rows) != released { 222 t.Fatalf("persisted rows %d do not match released count %d", len(store.rows), released) 223 } 224 if len(au.events) != released { 225 t.Fatalf("audit events %d do not match released count %d", len(au.events), released) 226 } 227} 228 229// TestGateFailedWriteIsNotSilentlyAudited proves a ledger insert failure surfaces as an error and the release 230// is NOT counted as audited beyond what the insert reached, so a write failure fails closed. 231func TestGateFailedWriteIsNotSilentlyAudited(t *testing.T) { 232 store := &fakeLedger{insertErr: errors.New("boom")} 233 au := &fakeAuditor{} 234 gate := newTestGate(store, au, 100, 1e-6) 235 236 rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}} 237 _, err := gate.Charge(context.Background(), rel) 238 if err == nil { 239 t.Fatalf("expected the ledger insert failure to surface") 240 } 241 if len(au.events) != 0 { 242 t.Fatalf("a failed ledger write was still audited: %d events", len(au.events)) 243 } 244} 245 246// TestGateAuditFailureFailsClosed proves a release whose audit cannot be written returns an error (fails 247// closed) rather than silently succeeding. 248func TestGateAuditFailureFailsClosed(t *testing.T) { 249 store := &fakeLedger{} 250 au := &fakeAuditor{failErr: errors.New("audit down")} 251 gate := newTestGate(store, au, 100, 1e-6) 252 253 rel := Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}} 254 _, err := gate.Charge(context.Background(), rel) 255 if err == nil { 256 t.Fatalf("expected the audit failure to fail the release closed") 257 } 258} 259 260// TestGateRemainingFlagOffIsZero proves Remaining returns zero headroom while the flag is off so nothing reads 261// a live budget on a disabled deploy. 262func TestGateRemainingFlagOffIsZero(t *testing.T) { 263 cfg := config.Federated{DPEnabled: false, DPEpsilon: 10, DPDelta: 1e-7} 264 gate := NewSpendGate(cfg, &fakeLedger{}, nil) 265 rem, err := gate.Remaining(context.Background(), "v1.4.0", nil) 266 if err != nil { 267 t.Fatalf("Remaining flag off: %v", err) 268 } 269 if rem != 0 { 270 t.Fatalf("flag off Remaining = %v, want 0", rem) 271 } 272} 273 274// TestAggregatorEmitRoutesThroughGate proves the active aggregator's Emit goes through the gate (an in budget 275// release persists and audits), confirming the Aggregator boundary is wired to the hard gate. 276func TestAggregatorEmitRoutesThroughGate(t *testing.T) { 277 store := &fakeLedger{} 278 au := &fakeAuditor{} 279 gate := newTestGate(store, au, 100, 1e-6) 280 agg := &activeAggregator{gate: gate} 281 282 _, err := agg.Emit(context.Background(), Release{ModelVersion: "v1.4.0", ContributingTenantCount: 5, Gaussian: GaussianRelease{NoiseSigma: 2, ClipNorm: 1}}) 283 if err != nil { 284 t.Fatalf("Emit through the gate: %v", err) 285 } 286 if len(store.rows) != 1 || len(au.events) != 1 { 287 t.Fatalf("Emit did not route through the gate: rows=%d events=%d", len(store.rows), len(au.events)) 288 } 289} 290