fix mocks in unit tests

Author: charleslien

.github/workflows/ci.yml

@@ -82,15 +82,14 @@ jobs:
     needs: [build-and-check]
     strategy:
       matrix:
-        package:
-          [.agents, backend, cli, common, npm-app, packages/agent-runtime, web]
+        package: [.agents, backend, cli, common, npm-app, packages, web]
         include:
           - package: .agents
           - package: backend
           - package: cli
           - package: common
           - package: npm-app
-          - package: packages/agent-runtime
+          - package: packages
           - package: web
     name: test-${{ matrix.package }}
     runs-on: ubuntu-latest
@@ -156,15 +155,22 @@ jobs:
     needs: [build-and-check]
     strategy:
       matrix:
-        package:
-          [.agents, backend, cli, common, npm-app, packages/agent-runtime, web]
+        package: [
+            .agents,
+            backend,
+            cli,
+            common,
+            npm-app,
+            packages
+            web,
+          ]
         include:
           - package: .agents
           - package: backend
           - package: cli
           - package: common
           - package: npm-app
-          - package: packages/agent-runtime
+          - package: packages
           - package: web
     name: test-integration-${{ matrix.package }}
     runs-on: ubuntu-latest

backend/src/__tests__/auto-topup.test.ts

@@ -45,7 +45,7 @@ describe('Auto Top-up System', () => {
       )
 
       // Mock the database
-      mockModule('@codebuff/common/db', () => ({
+      mockModule('@codebuff/internal/db', () => ({
         default: {
           query: {
             user: {

backend/src/__tests__/usage-calculation.test.ts

@@ -1,8 +1,8 @@
 import { calculateUsageAndBalance } from '@codebuff/billing'
 import { GRANT_PRIORITIES } from '@codebuff/common/constants/grant-priorities'
 import {
-  clearMockedModules,
   mockModule,
+  clearMockedModules,
 } from '@codebuff/common/testing/mock-modules'
 import { afterAll, beforeAll, describe, expect, it } from 'bun:test'
 
@@ -19,7 +19,7 @@ describe('Usage Calculation System', () => {
 
   beforeAll(() => {
     // Mock the database module before importing the function
-    mockModule('@codebuff/common/db', () => ({
+    mockModule('@codebuff/internal/db', () => ({
       default: {
         select: () => ({
           from: () => ({
@@ -61,7 +61,7 @@ describe('Usage Calculation System', () => {
     ]
 
     // Mock the database module with the test data
-    mockModule('@codebuff/common/db', () => ({
+    mockModule('@codebuff/internal/db', () => ({
       default: {
         select: () => ({
           from: () => ({
@@ -98,7 +98,7 @@ describe('Usage Calculation System', () => {
     ]
 
     // Mock the database module with the test data
-    mockModule('@codebuff/common/db', () => ({
+    mockModule('@codebuff/internal/db', () => ({
       default: {
         select: () => ({
           from: () => ({
@@ -145,7 +145,7 @@ describe('Usage Calculation System', () => {
     ]
 
     // Mock the database module with the test data
-    mockModule('@codebuff/common/db', () => ({
+    mockModule('@codebuff/internal/db', () => ({
       default: {
         select: () => ({
           from: () => ({
@@ -201,7 +201,7 @@ describe('Usage Calculation System', () => {
     ]
 
     // Mock the database module with the test data
-    mockModule('@codebuff/common/db', () => ({
+    mockModule('@codebuff/internal/db', () => ({
       default: {
         select: () => ({
           from: () => ({

packages/agent-runtime/src/templates/__tests__/agent-registry.test.ts

@@ -1,5 +1,4 @@
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
-import { mockModule } from '@codebuff/common/testing/mock-modules'
 import { getStubProjectFileContext } from '@codebuff/common/util/file'
 import {
   describe,
@@ -9,7 +8,6 @@ import {
   afterEach,
   spyOn,
   mock,
-  beforeAll,
 } from 'bun:test'
 
 import {
@@ -69,44 +67,6 @@ const mockStaticTemplates: Record<string, AgentTemplate> = {
 describe('Agent Registry', () => {
   let mockFileContext: ProjectFileContext
 
-  beforeAll(() => {
-    // Mock the database module
-    mockModule('@codebuff/common/db', () => ({
-      default: {
-        select: () => ({
-          from: () => ({
-            where: () => ({
-              orderBy: () => ({
-                limit: () => Promise.resolve([]),
-              }),
-              then: (fn: (rows: any[]) => any) => fn([]),
-            }),
-          }),
-        }),
-      },
-    }))
-
-    // Mock the schema module
-    mockModule('@codebuff/common/db/schema', () => ({
-      agentConfig: {
-        id: 'id',
-        publisher_id: 'publisher_id',
-        version: 'version',
-        major: 'major',
-        minor: 'minor',
-        patch: 'patch',
-        data: 'data',
-      },
-    }))
-
-    // Mock drizzle-orm
-    mockModule('drizzle-orm', () => ({
-      and: (...args: any[]) => ({ type: 'and', args }),
-      desc: (field: any) => ({ type: 'desc', field }),
-      eq: (field: any, value: any) => ({ type: 'eq', field, value }),
-    }))
-  })
-
   beforeEach(async () => {
     agentRuntimeImpl = {
       ...TEST_AGENT_RUNTIME_IMPL,

packages/billing/src/__tests__/credit-delegation.test.ts

@@ -31,7 +31,7 @@ describe('Credit Delegation', () => {
     }))
 
     // Mock common dependencies
-    await mockModule('@codebuff/common/db', () => {
+    await mockModule('@codebuff/internal/db', () => {
       const select = mock((fields: Record<string, unknown>) => {
         if ('orgId' in fields && 'orgName' in fields) {
           return {
@@ -78,17 +78,6 @@ describe('Credit Delegation', () => {
         },
       }
     })
-
-    await mockModule('@codebuff/common/db/schema', () => ({
-      orgMember: { org_id: 'org_id', user_id: 'user_id' },
-      org: { id: 'id', name: 'name', slug: 'slug' },
-      orgRepo: {
-        org_id: 'org_id',
-        repo_url: 'repo_url',
-        repo_name: 'repo_name',
-        is_active: 'is_active',
-      },
-    }))
   })
 
   afterAll(() => {

packages/billing/src/__tests__/org-billing.test.ts

@@ -2,7 +2,7 @@ import {
   clearMockedModules,
   mockModule,
 } from '@codebuff/common/testing/mock-modules'
-import { afterAll, afterEach, beforeAll, describe, expect, it } from 'bun:test'
+import { afterEach, beforeEach, describe, expect, it } from 'bun:test'
 
 import {
   calculateOrganizationUsageAndBalance,
@@ -80,12 +80,11 @@ const createDbMock = (options?: {
 }
 
 describe('Organization Billing', () => {
-  beforeAll(async () => {
-    await mockModule('@codebuff/common/db', () => ({
+  beforeEach(async () => {
+    await mockModule('@codebuff/internal/db', () => ({
       default: createDbMock(),
     }))
-
-    await mockModule('@codebuff/common/db/transaction', () => ({
+    await mockModule('@codebuff/internal/db/transaction', () => ({
       withSerializableTransaction: async ({
         callback,
       }: {
@@ -94,16 +93,10 @@ describe('Organization Billing', () => {
     }))
   })
 
-  afterAll(() => {
+  afterEach(() => {
     clearMockedModules()
   })
 
-  afterEach(async () => {
-    await mockModule('@codebuff/common/db', () => ({
-      default: createDbMock(),
-    }))
-  })
-
   describe('calculateOrganizationUsageAndBalance', () => {
     it('should calculate balance correctly with positive and negative balances', async () => {
       const organizationId = 'org-123'
@@ -130,7 +123,7 @@ describe('Organization Billing', () => {
 
     it('should handle organization with no grants', async () => {
       // Mock empty grants
-      await mockModule('@codebuff/common/db', () => ({
+      await mockModule('@codebuff/internal/db', () => ({
         default: createDbMock({ grants: [] }),
       }))
 
@@ -258,7 +251,7 @@ describe('Organization Billing', () => {
 
     it('should handle duplicate operation IDs gracefully', async () => {
       // Mock database constraint error
-      await mockModule('@codebuff/common/db', () => ({
+      await mockModule('@codebuff/internal/db', () => ({
         default: createDbMock({
           insert: () => ({
             values: () => {

packages/internal/src/utils/__tests__/version-utils.test.ts

@@ -128,7 +128,7 @@ describe('version-utils', () => {
   describe('getLatestAgentVersion', () => {
     it('should return version 0.0.0 when no agent exists', async () => {
       // Mock the database to return empty result
-      mockModule('@codebuff/common/db', () => ({
+      mockModule('@codebuff/internal/db', () => ({
         default: {
           select: () => ({
             from: () => ({
@@ -150,7 +150,7 @@ describe('version-utils', () => {
 
     it('should return latest version when agent exists', async () => {
       // Mock the database to return a version
-      mockModule('@codebuff/common/db', () => ({
+      mockModule('@codebuff/internal/db', () => ({
         default: {
           select: () => ({
             from: () => ({
@@ -173,7 +173,7 @@ describe('version-utils', () => {
 
     it('should handle null values in database response', async () => {
       // Mock the database to return null values
-      mockModule('@codebuff/common/db', () => ({
+      mockModule('@codebuff/internal/db', () => ({
         default: {
           select: () => ({
             from: () => ({
@@ -268,7 +268,7 @@ describe('version-utils', () => {
   describe('versionExists', () => {
     it('should return true when version exists', async () => {
       // Mock the database to return a result
-      mockModule('@codebuff/common/db', () => ({
+      mockModule('@codebuff/internal/db', () => ({
         default: {
           select: () => ({
             from: () => ({
@@ -290,7 +290,7 @@ describe('version-utils', () => {
 
     it('should return false when version does not exist', async () => {
       // Mock the database to return empty result
-      mockModule('@codebuff/common/db', () => ({
+      mockModule('@codebuff/internal/db', () => ({
         default: {
           select: () => ({
             from: () => ({