refactor: migrate billing and test files to dependency injection pattern #408
+6,785
−865
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- 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.
…ation Replaced outdated "Mock database module directly" guidance with: - Complete table of all DI interfaces and injectable dependencies - Code examples showing how to use deps parameters - Testing best practices for billing functions - Example test demonstrating the DI pattern This helps future developers understand how to properly test billing functions using the DI patterns established in this PR.
brandonkachen
force-pushed
the
billing-di-refactor-v2
branch
from
e9dc782 to
80d5a92
Compare
- Add Clickable wrapper component for interactive areas with non-selectable text
- Export makeTextUnselectable utility for edge cases
- Refactor Button to use shared makeTextUnselectable from Clickable
- Remove unused textToCopy prop from CopyButton
- Add selectable={false} to copy button text elements in message-block.tsx
- Document Button/Clickable usage patterns in cli/knowledge.md
… docs - Add CodebuffTransaction and CodebuffTransactionFn types to @codebuff/internal/db for fully-typed Drizzle transaction usage in production code - Improve BillingTransactionFn documentation explaining why any is needed (Drizzle PgTransaction signatures incompatible with BillingDbConnection) - Document when to use CodebuffTransactionFn vs BillingTransactionFn
…r types - error.ts: use unknown instead of any for error parameters - promise.ts: use unknown for error callbacks with proper type narrowing - object.ts: use Record<string, unknown> and proper generic types - split-data.ts: add JsonValue type, use unknown for public API - org-billing.ts: document WithSerializableTransactionFn any usage
…ck-db limitations
- Improve mock-db.ts: add filtering support to findFirst, document field-based query detection limitations - Remove type assertions in auto-topup.ts by conditionally skipping disableAutoTopupInternal when using DI - Fix WithSerializableTransactionFn duplication in org-billing.ts - Replace any types with Record<string, unknown> in test mocks - Add detailed JSDoc for shouldAttemptStripeMetering explaining CI=true behavior - Add comments explaining why DbConn is defined locally (Drizzle type compatibility)
…ency interfaces Addresses Gemini CLI review feedback: - balance-calculator.ts: CalculateUsageThisCycleDeps now uses BillingDbConnection - credit-delegation.ts: FindOrganizationForRepositoryDeps now uses BillingDbConnection - grant-credits.ts: GetPreviousFreeGrantAmountDeps and CalculateTotalReferralBonusDeps now use BillingDbConnection - org-billing.ts: SyncOrganizationBillingCycleDeps now uses BillingDbConnection Added type casts to resolve union type incompatibilities between real Drizzle db and mock BillingDbConnection interface.
- Replace `as unknown as` double-casts with cleaner alternatives: - balance-calculator.ts: use .then() handler for type assertion - credit-delegation.ts: use .then() handler for type assertion - grant-credits.ts: use direct type assertion (limit() returns sync) - All billing files now use consistent, cleaner type patterns
…al consistency
- Add centralized getBillingDbClient() helper to avoid repeated casting
- Use generic type parameters on .select<T>() for cleaner type inference
- Fix transactional consistency bug: pass deps: { db: tx } to
getPreviousFreeGrantAmount and calculateTotalReferralBonus in
triggerMonthlyResetAndGrant so they use the transaction instead of global db
Addresses all issues identified by Codex, Claude Code, and Gemini CLI reviews.
Apply the same getBillingDbClient helper pattern from grant-credits.ts to: - balance-calculator.ts (calculateUsageThisCycle) - credit-delegation.ts (findOrganizationForRepository) - org-billing.ts (syncOrganizationBillingCycle) This centralizes the DB client casting logic and removes redundant comments.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Non-Test Files Changed (22 files, +1,393 / -157 lines)
TESTING.mdcommon/src/types/contracts/billing.tspackages/billing/src/auto-topup.tspackages/billing/src/grant-credits.tspackages/billing/src/billing.knowledge.mdpackages/billing/src/org-billing.tspackages/billing/src/stripe-metering.tspackages/billing/src/balance-calculator.tspackages/internal/src/db/types.tspackages/billing/src/usage-service.tspackages/billing/src/credit-delegation.tspackages/agent-runtime/src/llm-api/linkup-api.tsweb/src/app/api/admin/relabel-for-user/route.tspackages/billing/src/org-monitoring.tsevals/buffbench/main*.ts(4 files)packages/internal/src/db/index.tscommon/src/old-constants.tsweb/src/app/api/v1/agent-runs/[runId]/steps/_post.tsweb/src/app/api/v1/agent-runs/_post.tsSummary
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
BillingDbConnection,UsageServiceDeps, etc.)createMockDb,createMockTransaction,createTrackedMockDb)createMockCreditGrant,createMockUser, etc.)Billing Package Refactor
grant-credits.tsto accept optionaldepswithtransactionfor DIusage-service.tsto acceptUsageServiceDepsfor testingmockModulegetBillingDbClient()helper for centralized DB client handlingtriggerMonthlyResetAndGrantTEST_USER_ID Removal
TEST_USER_IDbypass from production billing codebalance-calculator.ts,stripe-metering.ts, and agent-runs API endpointsCI=trueat entry points to skip Stripe billingTest Coverage
mockModuleusage remains in billing testsDocumentation
TESTING.mdwith comprehensive DI patterns and test fixture documentationbilling.knowledge.mdwith DI patterns referenceValidation
mockModuleusage remains in billing tests