From 2ec7f5c69b955783801c955be8dc5557f1dae678 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Mon, 6 Jul 2026 14:18:02 -0400 Subject: [PATCH] test: fix TestUserAIBudgetOverride audit flake (#27013) Closes https://linear.app/codercom/issue/AIGOV-502/flake-testuseraibudgetoverrideauditupsertspendlimit The `Audit/UpsertEverything` and `Audit/UpsertSpendLimit` subtests grabbed `rows[0]` from `GetAuditLogsOffset`, assuming strict time-desc order. Both upserts emit `AuditActionWrite` against the same override with no ordering tiebreaker, so when the create and update entries share a timestamp the create entry could sort first and fail the assertion. Select the update entry by the spend limit it results in via a new `auditLogByNewSpendLimit` helper instead of by row position.
We could make `auditLogByNewSpendLimit` more generic, but since we don't have any other use cases for it right now, I think we can postpone that until the need arises. Related PR: https://github.com/coder/coder/pull/26630/changes --- enterprise/coderd/aibridge_test.go | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 9910fb128f..52540bf0c7 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -59,6 +59,24 @@ func auditLogsByAction(t *testing.T, rows []database.GetAuditLogsOffsetRow) map[ return byAction } +// auditLogByNewSpendLimit selects the audit row whose diff sets spend_limit to +// the given value so callers can assert on a specific entry without relying on +// row order, which GetAuditLogsOffset does not guarantee. It requires the +// resulting spend_limit among the rows to be unique. +func auditLogByNewSpendLimit(t *testing.T, rows []database.GetAuditLogsOffsetRow, newSpendLimit string) database.AuditLog { + t.Helper() + var matches []database.AuditLog + for _, r := range rows { + var diff audit.Map + require.NoError(t, json.Unmarshal(r.AuditLog.Diff, &diff)) + if field, ok := diff["spend_limit"]; ok && field.New == newSpendLimit { + matches = append(matches, r.AuditLog) + } + } + require.Lenf(t, matches, 1, "want exactly one audit entry setting spend_limit to %q, got %d", newSpendLimit, len(matches)) + return matches[0] +} + func TestAIBridgeListSessions(t *testing.T) { t.Parallel() @@ -2687,8 +2705,10 @@ func TestUserAIBudgetOverride(t *testing.T) { ) require.NoError(t, err) require.Len(t, rows, 2, "expected one create and one update audit entry") - // GetAuditLogsOffset returns entries sorted by time in descending order. - updateLog := rows[0].AuditLog + // Both upserts emit AuditActionWrite; select the update by the spend + // limit it results in rather than row order, which GetAuditLogsOffset + // does not guarantee. + updateLog := auditLogByNewSpendLimit(t, rows, "$1000.00") var updateDiff audit.Map require.NoError(t, json.Unmarshal(updateLog.Diff, &updateDiff)) @@ -2744,8 +2764,10 @@ func TestUserAIBudgetOverride(t *testing.T) { ) require.NoError(t, err) require.Len(t, rows, 2, "expected one create and one update audit entry") - // GetAuditLogsOffset returns entries sorted by time in descending order. - updateLog := rows[0].AuditLog + // Both upserts emit AuditActionWrite; select the update by the spend + // limit it results in rather than row order, which GetAuditLogsOffset + // does not guarantee. + updateLog := auditLogByNewSpendLimit(t, rows, "$1000.00") var updateDiff audit.Map require.NoError(t, json.Unmarshal(updateLog.Diff, &updateDiff))