Merge branch 'main' into set_output

Author: jahooma

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

@@ -1,13 +1,17 @@
-import type { CreditBalance } from '@codebuff/billing'
-import { checkAndTriggerAutoTopup } from '@codebuff/billing'
 import * as billing from '@codebuff/billing'
+import { checkAndTriggerAutoTopup } from '@codebuff/billing'
 import {
+  clearMockedModules,
+  mockModule,
+} from '@codebuff/common/testing/mock-modules'
+import {
+  afterAll,
+  afterEach,
   beforeEach,
   describe,
   expect,
   it,
   mock,
-  afterEach,
   spyOn,
 } from 'bun:test'
 
@@ -60,7 +64,7 @@ describe('Auto Top-up System', () => {
       grantCreditsMock = mock(() => Promise.resolve())
 
       // Mock the database
-      mock.module('@codebuff/common/db', () => ({
+      mockModule('@codebuff/common/db', () => ({
         default: {
           query: {
             user: {
@@ -84,7 +88,7 @@ describe('Auto Top-up System', () => {
       )
 
       // Mock Stripe payment intent creation
-      mock.module('@codebuff/common/util/stripe', () => ({
+      mockModule('@codebuff/common/util/stripe', () => ({
         stripeServer: {
           paymentIntents: {
             create: mock(() =>
@@ -98,6 +102,10 @@ describe('Auto Top-up System', () => {
       }))
     })
 
+    afterAll(() => {
+      clearMockedModules()
+    })
+
     it('should trigger top-up when balance below threshold', async () => {
       await checkAndTriggerAutoTopup('test-user')
 

backend/src/__tests__/dynamic-agent-loader.test.ts

@@ -1,35 +1,14 @@
+import {
+  clearMockedModules,
+  mockModule,
+} from '@codebuff/common/testing/mock-modules'
 import { ProjectFileContext } from '@codebuff/common/util/file'
-import { describe, expect, it, mock } from 'bun:test'
+import { afterAll, beforeAll, describe, expect, it } from 'bun:test'
 import {
   dynamicAgentService,
   DynamicAgentService,
 } from '../templates/dynamic-agent-service'
 
-// Mock backend utility module
-mock.module('../util/file-resolver', () => ({
-  resolvePromptField: (field: string | { path: string }, basePath: string) => {
-    if (typeof field === 'string') {
-      return field
-    }
-    if (field.path?.includes('brainstormer-system.md')) {
-      return 'You are a creative brainstormer.'
-    }
-    if (field.path?.includes('brainstormer-user-input.md')) {
-      return 'Help brainstorm ideas.'
-    }
-    return 'Mock content'
-  },
-  resolveFileContent: (filePath: string, basePath: string) => {
-    if (filePath.includes('brainstormer-system.md')) {
-      return 'You are a creative brainstormer.'
-    }
-    if (filePath.includes('brainstormer-user-input.md')) {
-      return 'Help brainstorm ideas.'
-    }
-    return 'Mock content'
-  },
-}))
-
 describe('Dynamic Agent Loader', () => {
   const mockFileContext: ProjectFileContext = {
     projectRoot: '/test',
@@ -56,6 +35,40 @@ describe('Dynamic Agent Loader', () => {
     },
   }
 
+  beforeAll(() => {
+    // Mock backend utility module
+    mockModule('@codebuff/backend/util/file-resolver', () => ({
+      resolvePromptField: (
+        field: string | { path: string },
+        basePath: string
+      ) => {
+        if (typeof field === 'string') {
+          return field
+        }
+        if (field.path?.includes('brainstormer-system.md')) {
+          return 'You are a creative brainstormer.'
+        }
+        if (field.path?.includes('brainstormer-user-input.md')) {
+          return 'Help brainstorm ideas.'
+        }
+        return 'Mock content'
+      },
+      resolveFileContent: (filePath: string, basePath: string) => {
+        if (filePath.includes('brainstormer-system.md')) {
+          return 'You are a creative brainstormer.'
+        }
+        if (filePath.includes('brainstormer-user-input.md')) {
+          return 'Help brainstorm ideas.'
+        }
+        return 'Mock content'
+      },
+    }))
+  })
+
+  afterAll(() => {
+    clearMockedModules()
+  })
+
   it('should load valid dynamic agent template', async () => {
     const fileContext: ProjectFileContext = {
       ...mockFileContext,

backend/src/__tests__/dynamic-agent-schema-validation.test.ts

@@ -1,28 +1,38 @@
+import {
+  clearMockedModules,
+  mockModule,
+} from '@codebuff/common/testing/mock-modules'
 import {
   getStubProjectFileContext,
   ProjectFileContext,
 } from '@codebuff/common/util/file'
-import { beforeEach, describe, expect, it, mock } from 'bun:test'
+import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'bun:test'
 import { DynamicAgentService } from '../templates/dynamic-agent-service'
 
-// Mock logger to avoid console output during tests
-mock.module('../util/logger', () => ({
-  logger: {
-    debug: () => {},
-    warn: () => {},
-    error: () => {},
-  },
-}))
-
 describe('Dynamic Agent Schema Validation', () => {
   let service: DynamicAgentService
   let mockFileContext: ProjectFileContext
 
+  beforeAll(() => {
+    // Mock logger to avoid console output during tests
+    mockModule('@codebuff/backend/util/logger', () => ({
+      logger: {
+        debug: () => {},
+        warn: () => {},
+        error: () => {},
+      },
+    }))
+  })
+
   beforeEach(() => {
     service = new DynamicAgentService()
     mockFileContext = getStubProjectFileContext()
   })
 
+  afterAll(() => {
+    clearMockedModules()
+  })
+
   describe('Default Schema Behavior', () => {
     it('should have no prompt schema when no promptSchema provided', async () => {
       const fileContext: ProjectFileContext = {

backend/src/__tests__/fast-rewrite.test.ts

@@ -1,30 +1,40 @@
-import { describe, expect, it, mock } from 'bun:test'
 import { TEST_USER_ID } from '@codebuff/common/constants'
+import {
+  clearMockedModules,
+  mockModule,
+} from '@codebuff/common/testing/mock-modules'
+import { afterAll, beforeAll, describe, expect, it } from 'bun:test'
 import { createPatch } from 'diff'
 import path from 'path'
 import { rewriteWithOpenAI } from '../fast-rewrite'
 
-// Mock database interactions
-mock.module('pg-pool', () => ({
-  Pool: class {
-    connect() {
-      return {
-        query: () => ({
-          rows: [{ id: 'test-user-id' }],
-          rowCount: 1,
-        }),
-        release: () => {},
-      }
-    }
-  },
-}))
+describe.skip('rewriteWithOpenAI', () => {
+  beforeAll(() => {
+    // Mock database interactions
+    mockModule('pg-pool', () => ({
+      Pool: class {
+        connect() {
+          return {
+            query: () => ({
+              rows: [{ id: 'test-user-id' }],
+              rowCount: 1,
+            }),
+            release: () => {},
+          }
+        }
+      },
+    }))
 
-// Mock message saving
-mock.module('backend/llm-apis/message-cost-tracker', () => ({
-  saveMessage: () => Promise.resolve(),
-}))
+    // Mock message saving
+    mockModule('@codebuff/backend/llm-apis/message-cost-tracker', () => ({
+      saveMessage: () => Promise.resolve(),
+    }))
+  })
+
+  afterAll(() => {
+    clearMockedModules()
+  })
 
-describe.skip('rewriteWithOpenAI', () => {
   it('should correctly integrate edit snippet changes while preserving formatting', async () => {
     const testDataDir = path.join(__dirname, 'test-data', 'dex-go')
     const originalContent = await Bun.file(`${testDataDir}/original.go`).text()

backend/src/__tests__/main-prompt.test.ts

@@ -3,7 +3,9 @@ import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/constants'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
+  afterAll,
   afterEach,
+  beforeAll,
   beforeEach,
   describe,
   expect,
@@ -25,28 +27,34 @@ import {
   getToolCallString,
   renderToolResults,
 } from '@codebuff/common/constants/tools'
+import {
+  clearMockedModules,
+  mockModule,
+} from '@codebuff/common/testing/mock-modules'
 import { ProjectFileContext } from '@codebuff/common/util/file'
 import * as getDocumentationForQueryModule from '../get-documentation-for-query'
 import * as websocketAction from '../websockets/websocket-action'
 
-// Mock logger
-mock.module('../util/logger', () => ({
-  logger: {
-    debug: () => {},
-    error: () => {},
-    info: () => {},
-    warn: () => {},
-  },
-  withLoggerContext: async (context: any, fn: () => Promise<any>) => fn(),
-}))
-
 const mockAgentStream = (streamOutput: string) => {
   spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* () {
     yield streamOutput
   })
 }
 
 describe('mainPrompt', () => {
+  beforeAll(() => {
+    // Mock logger
+    mockModule('@codebuff/backend/util/logger', () => ({
+      logger: {
+        debug: () => {},
+        error: () => {},
+        info: () => {},
+        warn: () => {},
+      },
+      withLoggerContext: async (context: any, fn: () => Promise<any>) => fn(),
+    }))
+  })
+
   beforeEach(() => {
     // Mock analytics and tracing
     spyOn(analytics, 'initAnalytics').mockImplementation(() => {})
@@ -141,6 +149,10 @@ describe('mainPrompt', () => {
     mock.restore()
   })
 
+  afterAll(() => {
+    clearMockedModules()
+  })
+
   class MockWebSocket {
     send(msg: string) {}
     close() {}