fix: address PR review feedback

brandonkachen · brandonkachen · commit 0ab7a1d65a8c · 2026-01-12T15:42:11.000-08:00
- 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
diff --git a/packages/billing/src/__tests__/grant-credits.test.ts b/packages/billing/src/__tests__/grant-credits.test.ts
@@ -289,18 +289,30 @@ describe('grant-credits', () => {
               where: () => Promise.resolve(),
             }),
           }),
-          // Note: grantCreditOperation uses `db.insert` directly when no debt,
-          // not `tx.insert`. This test verifies the debt-checking path works.
+          insert: () => ({
+            values: (values: any) => {
+              insertedGrants.push(values)
+              return Promise.resolve()
+            },
+          }),
         }
 
-        // Since the function uses db.insert directly for no-debt case,
-        // we test by verifying that the debt check path returns empty array
-        // and the function completes without error
-        // The actual insert would hit the real DB, so we verify the flow instead
-        
-        // For a pure unit test, we'd need to inject the db dependency too
-        // This is a limitation of the current DI pattern
-        expect(true).toBe(true) // Placeholder - function flow verified
+        await grantCreditOperation({
+          userId: 'user-123',
+          amount: 500,
+          type: 'free',
+          description: 'Monthly free credits',
+          expiresAt: futureDate,
+          operationId: 'new-grant-1',
+          tx: mockTx as any,
+          logger,
+        })
+
+        // Should have created a new grant with full balance (no debt to deduct)
+        expect(insertedGrants.length).toBe(1)
+        expect(insertedGrants[0].principal).toBe(500)
+        expect(insertedGrants[0].balance).toBe(500)
+        expect(insertedGrants[0].description).toBe('Monthly free credits')
       })
 
       it('should not create grant when debt exceeds amount', async () => {
@@ -368,165 +380,149 @@ describe('grant-credits', () => {
   describe('revokeGrantByOperationId', () => {
     it('should successfully revoke a grant with positive balance', async () => {
       const updatedValues: any[] = []
-      let transactionCalled = false
-
-      // Mock db.transaction
-      const mockDb = {
-        transaction: async <T>(callback: (tx: any) => Promise<T>): Promise<T> => {
-          transactionCalled = true
-          const tx = {
-            query: {
-              creditLedger: {
-                findFirst: async () => ({
-                  operation_id: 'grant-to-revoke',
-                  user_id: 'user-123',
-                  principal: 500,
-                  balance: 300, // 200 already consumed
-                  type: 'purchase',
-                  description: 'Purchased 500 credits',
-                }),
-              },
-            },
-            update: () => ({
-              set: (values: any) => ({
-                where: () => {
-                  updatedValues.push(values)
-                  return Promise.resolve()
-                },
+
+      const mockTransaction: BillingTransactionFn = async (callback) => {
+        const tx = {
+          query: {
+            creditLedger: {
+              findFirst: async () => ({
+                operation_id: 'grant-to-revoke',
+                user_id: 'user-123',
+                principal: 500,
+                balance: 300, // 200 already consumed
+                type: 'purchase',
+                description: 'Purchased 500 credits',
               }),
+            },
+          },
+          update: () => ({
+            set: (values: any) => ({
+              where: () => {
+                updatedValues.push(values)
+                return Promise.resolve()
+              },
             }),
-          }
-          return callback(tx)
-        },
+          }),
+        }
+        return callback(tx)
       }
 
-      // Use the actual function with mocked db
-      const result = await mockDb.transaction(async (tx) => {
-        const grant = await tx.query.creditLedger.findFirst({})
-        if (!grant) return false
-        if (grant.balance < 0) return false
-
-        await tx
-          .update()
-          .set({
-            principal: 0,
-            balance: 0,
-            description: `${grant.description} (Revoked: Test refund)`,
-          })
-          .where()
-
-        return true
+      const result = await revokeGrantByOperationId({
+        operationId: 'grant-to-revoke',
+        reason: 'Test refund',
+        logger,
+        deps: { transaction: mockTransaction },
       })
 
       expect(result).toBe(true)
-      expect(transactionCalled).toBe(true)
       expect(updatedValues.length).toBe(1)
       expect(updatedValues[0].principal).toBe(0)
       expect(updatedValues[0].balance).toBe(0)
       expect(updatedValues[0].description).toContain('Revoked: Test refund')
     })
 
     it('should return false when grant does not exist', async () => {
-      const mockDb = {
-        transaction: async <T>(callback: (tx: any) => Promise<T>): Promise<T> => {
-          const tx = {
-            query: {
-              creditLedger: {
-                findFirst: async () => null, // Grant not found
-              },
+      const mockTransaction: BillingTransactionFn = async (callback) => {
+        const tx = {
+          query: {
+            creditLedger: {
+              findFirst: async () => null, // Grant not found
             },
-            update: () => ({
-              set: () => ({
-                where: () => Promise.resolve(),
-              }),
+          },
+          update: () => ({
+            set: () => ({
+              where: () => Promise.resolve(),
             }),
-          }
-          return callback(tx)
-        },
+          }),
+        }
+        return callback(tx)
       }
 
-      const result = await mockDb.transaction(async (tx) => {
-        const grant = await tx.query.creditLedger.findFirst({})
-        if (!grant) return false
-        return true
+      const result = await revokeGrantByOperationId({
+        operationId: 'non-existent-grant',
+        reason: 'Test refund',
+        logger,
+        deps: { transaction: mockTransaction },
       })
 
       expect(result).toBe(false)
     })
 
     it('should return false when grant has negative balance', async () => {
-      const mockDb = {
-        transaction: async <T>(callback: (tx: any) => Promise<T>): Promise<T> => {
-          const tx = {
-            query: {
-              creditLedger: {
-                findFirst: async () => ({
-                  operation_id: 'debt-grant',
-                  user_id: 'user-123',
-                  principal: 500,
-                  balance: -100, // User has overspent
-                  type: 'free',
-                  description: 'Monthly free credits',
-                }),
-              },
-            },
-            update: () => ({
-              set: () => ({
-                where: () => Promise.resolve(),
+      const mockTransaction: BillingTransactionFn = async (callback) => {
+        const tx = {
+          query: {
+            creditLedger: {
+              findFirst: async () => ({
+                operation_id: 'debt-grant',
+                user_id: 'user-123',
+                principal: 500,
+                balance: -100, // User has overspent
+                type: 'free',
+                description: 'Monthly free credits',
               }),
+            },
+          },
+          update: () => ({
+            set: () => ({
+              where: () => Promise.resolve(),
             }),
-          }
-          return callback(tx)
-        },
+          }),
+        }
+        return callback(tx)
       }
 
-      const result = await mockDb.transaction(async (tx) => {
-        const grant = await tx.query.creditLedger.findFirst({})
-        if (!grant) return false
-        if (grant.balance < 0) return false // Cannot revoke debt
-        return true
+      const result = await revokeGrantByOperationId({
+        operationId: 'debt-grant',
+        reason: 'Test refund',
+        logger,
+        deps: { transaction: mockTransaction },
       })
 
       expect(result).toBe(false)
     })
 
-    it('should return false when grant balance is exactly zero', async () => {
-      const mockDb = {
-        transaction: async <T>(callback: (tx: any) => Promise<T>): Promise<T> => {
-          const tx = {
-            query: {
-              creditLedger: {
-                findFirst: async () => ({
-                  operation_id: 'depleted-grant',
-                  user_id: 'user-123',
-                  principal: 500,
-                  balance: 0, // Fully consumed
-                  type: 'free',
-                  description: 'Monthly free credits',
-                }),
-              },
-            },
-            update: () => ({
-              set: (values: any) => ({
-                where: () => Promise.resolve(),
+    it('should successfully revoke a grant with zero balance', async () => {
+      const updatedValues: any[] = []
+
+      const mockTransaction: BillingTransactionFn = async (callback) => {
+        const tx = {
+          query: {
+            creditLedger: {
+              findFirst: async () => ({
+                operation_id: 'depleted-grant',
+                user_id: 'user-123',
+                principal: 500,
+                balance: 0, // Fully consumed
+                type: 'free',
+                description: 'Monthly free credits',
               }),
+            },
+          },
+          update: () => ({
+            set: (values: any) => ({
+              where: () => {
+                updatedValues.push(values)
+                return Promise.resolve()
+              },
             }),
-          }
-          return callback(tx)
-        },
+          }),
+        }
+        return callback(tx)
       }
 
-      // Note: The actual revokeGrantByOperationId checks for balance < 0,
-      // so a balance of 0 would still be revoked (nothing to revoke though)
-      const result = await mockDb.transaction(async (tx) => {
-        const grant = await tx.query.creditLedger.findFirst({})
-        if (!grant) return false
-        if (grant.balance < 0) return false
-        // Balance of 0 is technically revocable but there's nothing to revoke
-        return true
+      // Balance of 0 is not negative, so it can be revoked (nothing to actually revoke though)
+      const result = await revokeGrantByOperationId({
+        operationId: 'depleted-grant',
+        reason: 'Test refund',
+        logger,
+        deps: { transaction: mockTransaction },
       })
 
       expect(result).toBe(true)
+      expect(updatedValues.length).toBe(1)
+      expect(updatedValues[0].principal).toBe(0)
+      expect(updatedValues[0].balance).toBe(0)
     })
   })
 })
diff --git a/packages/billing/src/grant-credits.ts b/packages/billing/src/grant-credits.ts
@@ -12,7 +12,10 @@ import { and, desc, eq, gt, isNull, lte, or, sql } from 'drizzle-orm'
 import { generateOperationIdTimestamp } from './utils'
 
 import type { Logger } from '@codebuff/common/types/contracts/logger'
-import type { TriggerMonthlyResetAndGrantDeps } from '@codebuff/common/types/contracts/billing'
+import type {
+  TriggerMonthlyResetAndGrantDeps,
+  BillingTransactionFn,
+} from '@codebuff/common/types/contracts/billing'
 import type { GrantType } from '@codebuff/internal/db/schema'
 
 type CreditGrantSelect = typeof schema.creditLedger.$inferSelect
@@ -201,7 +204,7 @@ export async function grantCreditOperation(params: {
   } else {
     // No debt - create grant normally
     try {
-      await db.insert(schema.creditLedger).values({
+      await dbClient.insert(schema.creditLedger).values({
         operation_id: operationId,
         user_id: userId,
         principal: amount,
@@ -290,16 +293,19 @@ export async function processAndGrantCredit(params: {
  *
  * @param operationId The operation ID of the grant to revoke
  * @param reason The reason for revoking the credits (e.g. refund)
+ * @param deps Optional dependencies for testing (transaction function)
  * @returns true if the grant was found and revoked, false otherwise
  */
 export async function revokeGrantByOperationId(params: {
   operationId: string
   reason: string
   logger: Logger
+  deps?: { transaction?: BillingTransactionFn }
 }): Promise<boolean> {
-  const { operationId, reason, logger } = params
+  const { operationId, reason, logger, deps = {} } = params
+  const transaction = deps.transaction ?? db.transaction.bind(db)
 
-  return await db.transaction(async (tx) => {
+  return await transaction(async (tx) => {
     const grant = await tx.query.creditLedger.findFirst({
       where: eq(schema.creditLedger.operation_id, operationId),
     })
diff --git a/packages/billing/src/org-billing.ts b/packages/billing/src/org-billing.ts
@@ -266,6 +266,10 @@ export async function calculateOrganizationUsageAndBalance(
 
 /**
  * Consumes credits from organization grants in priority order.
+ *
+ * TODO: Add DI support for testing. This function requires more complex mocking
+ * due to withSerializableTransaction usage. Better tested with integration tests
+ * that verify the full org credit flow works correctly.
  */
 export async function consumeOrganizationCredits(params: {
   organizationId: string
@@ -308,6 +312,10 @@ export async function consumeOrganizationCredits(params: {
 
 /**
  * Grants credits to an organization.
+ *
+ * TODO: Add DI support for testing. This function uses the global db directly.
+ * Better tested with integration tests that verify the full org credit flow
+ * works correctly with database constraints and idempotency checks.
  */
 export async function grantOrganizationCredits(
   params: OptionalFields<
