Skip to content

Conversation

@brandonkachen
Copy link
Collaborator

Summary

This PR refactors billing code and tests to use dependency injection (DI) instead of module mocking, following the patterns outlined in TESTING.md.

Changes

DI Infrastructure

  • Add DI contract types for billing dependencies (BillingDbConnection, UsageServiceDeps, etc.)
  • Create mock database helpers (createMockDb, createMockTransaction, createTrackedMockDb)
  • Create test fixtures for billing (createMockCreditGrant, createMockUser, etc.)

Billing Package Refactor

  • Refactor grant-credits.ts to accept optional deps with transaction for DI
  • Refactor usage-service.ts to accept UsageServiceDeps for testing
  • Update all billing tests to use DI instead of mockModule

TEST_USER_ID Removal

  • Remove hardcoded TEST_USER_ID bypass from production billing code
  • Remove from balance-calculator.ts, stripe-metering.ts, and agent-runs API endpoints
  • Evals now set CI=true at entry points to skip Stripe billing

Test Coverage

  • Add 28 new unit tests for billing DI patterns
  • Add tests for grantCreditOperation debt settlement
  • Add tests for revokeGrantByOperationId scenarios
  • Add tests for usage-service balance calculations and error handling
  • Add tests for credit-delegation URL utilities

Documentation

  • Add TESTING.md with comprehensive DI patterns and test fixture documentation
  • Update knowledge.md with reference to testing documentation

Results

  • ✅ All 73 billing tests pass
  • ✅ All 13 packages pass typecheck
  • ✅ Zero mockModule usage remains in billing tests

- Add DI contract types for billing dependencies (BillingDbConnection, UsageServiceDeps, etc.)
- Create mock database helpers (createMockDb, createMockTransaction, createTrackedMockDb)
- Create test fixtures for billing (TEST_USER_ID, createMockCreditGrant, createMockUser, etc.)
- Refactor billing package (grant-credits.ts, usage-service.ts) to accept optional deps
- Update all billing tests to use DI instead of mockModule
- Update 16 agent-runtime tests to import TEST_USER_ID from test fixtures
- Refactor linkup-api.ts to support DI for withTimeout
- Refactor code-search.ts to support DI for spawn
- Refactor ban-conditions.ts to support DI for db and stripeServer
- Refactor credentials-storage to use getConfigDirFromEnvironment for DI
- Remove all mockModule usage from test files (replaced with DI or spyOn)
- Add comprehensive TESTING.md documentation for DI patterns and test fixtures
- Update knowledge.md with reference to TESTING.md
- Remove TEST_USER_ID checks from balance-calculator.ts (billing bypass)
- Remove TEST_USER_ID checks from stripe-metering.ts (already has CI check)
- Remove TEST_USER_ID checks from agent-runs API endpoints
- Replace TEST_USER_ID with local ADMIN_RELABEL_USER_ID in admin route
- Remove TEST_USER_ID export from old-constants.ts
- Remove TEST_USER_ID from test fixtures (use inline strings instead)
- Update 18 test files to use inline test-user-id string

The hardcoded test user bypass was a hack that coupled test infrastructure
to production billing code. Tests now use inline strings and production
code relies on environment checks (CI=true) to skip billing in test envs.
These tests were testing the old TEST_USER_ID skip behavior that was
removed in the previous commit. Tests now go through normal flow.
Instead of adding a new SKIP_BILLING env var, evals now set CI=true at the
top of their entry points (main.ts, main-nightly.ts, main-single-eval.ts,
main-hard-tasks.ts). This ensures billing is always skipped when running
evals, without requiring users to remember to set an env var.

The shouldAttemptStripeMetering() function already checks CI=true, so this
approach is simpler and follows the existing pattern.
- Add grantCreditOperation debt settlement tests
- Add revokeGrantByOperationId scenario tests
- Add usage-service balance calculation tests with multiple grant types
- Add usage-service error handling tests
- Add credit-delegation normalizeRepositoryUrl and extractOwnerAndRepo tests
- Add credit-delegation URL edge case tests

73 billing tests now pass with full DI coverage.
The TEST_USER_ID constant was removed from fixtures, so update the
documentation examples to use inline test-user-id strings instead.
- Add monthly reset flow tests (future date skip, user not found error)
- Add usage data flow tests (reset -> auto-topup -> balance calculation)
- Add debt settlement flow tests (single/multiple debts, debt exceeds grant)
- Add credit delegation flow tests (no repo, empty repo, malformed URLs)
- Add credit fallback flow tests (no repo fallback, org delegation attempt)
- Add complete billing cycle tests (reset -> grant -> consume -> balance)
- Add balance calculation tests (multi-grant types, debt, personal context)
- Add error handling tests (user not found, DB errors, auto-topup graceful fail)

94 billing tests now pass with full DI coverage.
- Fix grantCreditOperation transaction bug: use dbClient instead of db
  when tx is provided to maintain transaction boundary
- Remove placeholder test and replace with real assertion
- Add DI support to revokeGrantByOperationId so tests can call the
  actual function with injected transaction
- Update revokeGrantByOperationId tests to call real function with DI
- Add TODO comments for consumeOrganizationCredits and
  grantOrganizationCredits noting they need integration tests
- Add DI support to consumeOrganizationCredits via deps.withSerializableTransaction
- Add DI support to grantOrganizationCredits via deps.db
- Remove TODO comments that were added in previous commit
- Add 11 integration tests for consumeOrganizationCredits:
  - Priority order consumption
  - Multi-grant consumption
  - No grants error handling
  - Purchased credits tracking
- Add 6 integration tests for grantOrganizationCredits:
  - Correct values creation
  - Default description
  - Idempotency (duplicate handling)
  - Error propagation
  - Priority setting
  - Expiration date

104 billing tests now pass with full DI coverage.
- Add typed query builder interfaces (SelectQueryBuilder, InsertQueryBuilder,
  UpdateQueryBuilder) with proper method chaining types
- Add WhereResult, FromResult, OrderByResult, GroupByResult interfaces
- Add TableQuery interface for findFirst operations with typed params
- Update BillingDbConnection to use generics instead of `any`
- Add JSDoc comments with usage examples throughout
- Add OnInsertCallback, OnUpdateCallback types for tracking operations
- Add TrackedMockDbResult interface for better autocomplete
- Convert MockDbConfig from type to interface with documentation

All 104 billing tests pass with improved type safety.
The callback parameter must use `any` because the real Drizzle transaction
type (PgTransaction) has many additional properties (schema, rollback, etc.)
that our minimal BillingDbConnection does not include. Using `any` allows
both the real transaction and mock implementations to work together.
- Change error: any to error: unknown in catch blocks with type guards
- Update isUniqueConstraintError helper to handle unknown type safely
- Change logContext/logData from any to Record<string, unknown>
- Improve org-monitoring.ts to avoid delete operator with destructuring

Files improved:
- grant-credits.ts: 4 error catch blocks
- stripe-metering.ts: 2 retry callbacks
- org-billing.ts: 1 error catch block + 2 log contexts
- auto-topup.ts: 1 log context
- org-monitoring.ts: 1 log object
The billing DI refactor is complete and validated with E2E tests.
Removing the planning documents as they are no longer needed.
Revert changes to:
- cli/src/utils/auth.ts (getConfigDirFromEnvironment extraction)
- sdk/src/tools/code-search.ts (spawn DI)
- web/src/lib/ban-conditions.ts (db/stripe DI)

These DI improvements can be done in a separate PR.
This PR now focuses solely on billing DI patterns.
Auto-topup functions:
- validateAutoTopupStatus: DI for db and stripeServer
- checkAndTriggerAutoTopup: DI for db, stripe, calculateUsageAndBalance, etc.
- checkAndTriggerOrgAutoTopup: DI for db, stripe, org billing functions

Stripe metering:
- reportPurchasedCreditsToStripe: DI for db, stripe, shouldAttemptStripeMetering

Grant credits helpers:
- getPreviousFreeGrantAmount: DI for db
- calculateTotalReferralBonus: DI for db
- processAndGrantCredit: DI for grantCreditFn and logSyncFailure

New test files:
- auto-topup.test.ts: 19 tests for auto-topup DI
- stripe-metering.test.ts: 12 tests for stripe metering DI
- Updated grant-credits.test.ts with helper function tests

Type improvements:
- Added BillingOrganization type to contracts/billing.ts
- Added org query to BillingDbConnection interface
- Added org support to mock-db.ts
Added dependency injection support to:
- syncOrganizationBillingCycle: DI for db and stripeServer
- findOrganizationForRepository: DI for db
- getOrderedActiveOrganizationGrants: already had conn param

New test file: org-billing.test.ts with 12 tests covering:
- Billing cycle sync error cases
- Stripe subscription syncing
- Organization grant ordering
- Repository lookup with DB mocking

All 180+ billing tests pass.
…ageThisCycle

Added dependency injection support to:
- consumeCreditsAndAddAgentStep: DI for withSerializableTransaction, trackEvent, reportPurchasedCreditsToStripe
- calculateUsageThisCycle: DI for db

New test file: balance-calculator.test.ts with 11 tests covering:
- Database query injection
- Transaction handling with mocked dependencies
- Analytics tracking verification
- Stripe reporting verification
- BYOK user handling
- Multi-grant consumption
- Error handling for missing grants and failed inserts
- Latency calculation

All 185+ billing tests pass.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants