Skip to content

Commit 25e41b0

Browse files
jahoomajahooma
committed
fix tests
1 parent 1cf3134 commit 25e41b0

File tree

6 files changed

+98
-142
lines changed

6 files changed

+98
-142
lines changed

‎backend/src/__tests__/cost-aggregation.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ describe('Cost Aggregation System', () => {
141141
messages: [],
142142
agentState: parentAgentState,
143143
sendSubagentChunk: () => {},
144+
system: 'Test system prompt',
144145
}
145146

146147
// Mock executeAgent to return results with different credit costs
@@ -213,6 +214,7 @@ describe('Cost Aggregation System', () => {
213214
messages: [],
214215
agentState: parentAgentState,
215216
sendSubagentChunk: () => {},
217+
system: 'Test system prompt',
216218
}
217219

218220
// Mock executeAgent to return success and failure with partial costs
@@ -369,6 +371,7 @@ describe('Cost Aggregation System', () => {
369371
messages: [],
370372
agentState: mainAgentState,
371373
sendSubagentChunk: () => {},
374+
system: 'Test system prompt',
372375
}
373376

374377
const mockExecuteAgent = spyOn(spawnAgentUtils, 'executeSubagent')

‎backend/src/__tests__/loop-agent-steps.test.ts‎

Lines changed: 44 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import {
1919
spyOn,
2020
} from 'bun:test'
2121

22+
import { withAppContext } from '../context/app-context'
2223
import { loopAgentSteps } from '../run-agent-step'
2324
import { clearAgentGeneratorCache } from '../run-programmatic-step'
2425
import { mockFileContext, MockWebSocket } from './test-utils'
@@ -35,6 +36,23 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
3536
let mockAgentState: AgentState
3637
let llmCallCount: number
3738

39+
const runLoopAgentStepsWithContext = async (
40+
ws: WebSocket,
41+
options: Parameters<typeof loopAgentSteps>[1],
42+
) => {
43+
return await withAppContext(
44+
{
45+
userId: options.userId,
46+
clientSessionId: options.clientSessionId,
47+
},
48+
{
49+
currentUserId: options.userId,
50+
processedRepoId: 'test-repo',
51+
},
52+
async () => loopAgentSteps(ws, options),
53+
)
54+
}
55+
3856
beforeAll(() => {
3957
// Mock logger
4058
mockModule('@codebuff/backend/util/logger', () => ({
@@ -64,9 +82,15 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
6482
getAgentPrompt: async () => 'Mock prompt',
6583
}))
6684

67-
// Mock live user inputs - will be overridden in individual tests
85+
// Mock live user inputs - default to true to allow tests to run
6886
mockModule('@codebuff/backend/live-user-inputs', () => ({
69-
checkLiveUserInput: () => false, // Default to false, override in tests
87+
checkLiveUserInput: () => true,
88+
resetLiveUserInputsState: () => {},
89+
startUserInput: () => {},
90+
endUserInput: () => {},
91+
cancelUserInput: () => {},
92+
setSessionConnected: () => {},
93+
getLiveUserInputIds: () => undefined,
7094
}))
7195

7296
// Mock file reading updates
@@ -87,6 +111,8 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
87111
})
88112

89113
beforeEach(() => {
114+
clearAgentGeneratorCache()
115+
90116
llmCallCount = 0
91117

92118
// Setup spies for database operations
@@ -156,8 +182,9 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
156182
})
157183

158184
afterEach(() => {
159-
mock.restore()
160185
clearAgentGeneratorCache()
186+
187+
mock.restore()
161188
})
162189

163190
afterAll(() => {
@@ -188,13 +215,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
188215
'test-agent': mockTemplate,
189216
}
190217

191-
// Mock checkLiveUserInput to allow the loop to continue
192-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
193-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
194-
() => true, // Always return true to allow loop to continue
195-
)
196-
197-
const result = await loopAgentSteps(
218+
const result = await runLoopAgentStepsWithContext(
198219
new MockWebSocket() as unknown as WebSocket,
199220
{
200221
userInputId: 'test-user-input',
@@ -241,7 +262,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
241262
'test-agent': mockTemplate,
242263
}
243264

244-
const result = await loopAgentSteps(
265+
const result = await runLoopAgentStepsWithContext(
245266
new MockWebSocket() as unknown as WebSocket,
246267
{
247268
userInputId: 'test-user-input',
@@ -290,13 +311,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
290311
'test-agent': mockTemplate,
291312
}
292313

293-
// Mock checkLiveUserInput to allow multiple iterations
294-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
295-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
296-
() => true, // Always return true to allow loop to continue
297-
)
298-
299-
const result = await loopAgentSteps(
314+
const result = await runLoopAgentStepsWithContext(
300315
new MockWebSocket() as unknown as WebSocket,
301316
{
302317
userInputId: 'test-user-input',
@@ -344,16 +359,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
344359
'test-agent': mockTemplate,
345360
}
346361

347-
let checkCallCount = 0
348-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
349-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
350-
() => {
351-
checkCallCount++
352-
return checkCallCount <= 5
353-
},
354-
)
355-
356-
const result = await loopAgentSteps(
362+
const result = await runLoopAgentStepsWithContext(
357363
new MockWebSocket() as unknown as WebSocket,
358364
{
359365
userInputId: 'test-user-input',
@@ -394,7 +400,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
394400
'test-agent': mockTemplate,
395401
}
396402

397-
const result = await loopAgentSteps(
403+
const result = await runLoopAgentStepsWithContext(
398404
new MockWebSocket() as unknown as WebSocket,
399405
{
400406
userInputId: 'test-user-input',
@@ -427,16 +433,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
427433
'test-agent': llmOnlyTemplate,
428434
}
429435

430-
let checkCallCount = 0
431-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
432-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
433-
() => {
434-
checkCallCount++
435-
return checkCallCount <= 2 // Allow 2 iterations
436-
},
437-
)
438-
439-
const result = await loopAgentSteps(
436+
const result = await runLoopAgentStepsWithContext(
440437
new MockWebSocket() as unknown as WebSocket,
441438
{
442439
userInputId: 'test-user-input',
@@ -471,16 +468,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
471468
'test-agent': mockTemplate,
472469
}
473470

474-
let checkCallCount = 0
475-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
476-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
477-
() => {
478-
checkCallCount++
479-
return checkCallCount <= 2
480-
},
481-
)
482-
483-
const result = await loopAgentSteps(
471+
const result = await runLoopAgentStepsWithContext(
484472
new MockWebSocket() as unknown as WebSocket,
485473
{
486474
userInputId: 'test-user-input',
@@ -532,16 +520,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
532520
'test-agent': mockTemplate,
533521
}
534522

535-
let checkCallCount = 0
536-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
537-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
538-
() => {
539-
checkCallCount++
540-
return checkCallCount <= 10 // Allow many iterations
541-
},
542-
)
543-
544-
const result = await loopAgentSteps(
523+
const result = await runLoopAgentStepsWithContext(
545524
new MockWebSocket() as unknown as WebSocket,
546525
{
547526
userInputId: 'test-user-input',
@@ -589,16 +568,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
589568
return getMessagesCallCount === 2 ? ['async message'] : []
590569
})
591570

592-
let checkCallCount = 0
593-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
594-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
595-
() => {
596-
checkCallCount++
597-
return checkCallCount <= 5
598-
},
599-
)
600-
601-
const result = await loopAgentSteps(
571+
const result = await runLoopAgentStepsWithContext(
602572
new MockWebSocket() as unknown as WebSocket,
603573
{
604574
userInputId: 'test-user-input',
@@ -649,13 +619,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
649619
'test-agent': mockTemplate,
650620
}
651621

652-
// Mock checkLiveUserInput to allow the loop to run
653-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
654-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
655-
() => true,
656-
)
657-
658-
await loopAgentSteps(new MockWebSocket() as unknown as WebSocket, {
622+
await runLoopAgentStepsWithContext(new MockWebSocket() as unknown as WebSocket, {
659623
userInputId: 'test-user-input',
660624
agentType: 'test-agent',
661625
agentState: mockAgentState,
@@ -737,20 +701,10 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
737701
return 'mock-message-id'
738702
})
739703

740-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
741-
let checkCount = 0
742-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
743-
() => {
744-
checkCount++
745-
return checkCount < 10 // Limit to prevent infinite loop
746-
},
747-
)
748-
749-
// Capture the agent state during execution
750704
mockAgentState.output = undefined
751705
capturedAgentState = mockAgentState
752706

753-
const result = await loopAgentSteps(
707+
const result = await runLoopAgentStepsWithContext(
754708
new MockWebSocket() as unknown as WebSocket,
755709
{
756710
userInputId: 'test-user-input',
@@ -821,15 +775,10 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
821775
return 'mock-message-id'
822776
})
823777

824-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
825-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
826-
() => true,
827-
)
828-
829778
mockAgentState.output = undefined
830779
capturedAgentState = mockAgentState
831780

832-
const result = await loopAgentSteps(
781+
const result = await runLoopAgentStepsWithContext(
833782
new MockWebSocket() as unknown as WebSocket,
834783
{
835784
userInputId: 'test-user-input',
@@ -876,12 +825,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
876825
return 'mock-message-id'
877826
})
878827

879-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
880-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
881-
() => true,
882-
)
883-
884-
const result = await loopAgentSteps(
828+
const result = await runLoopAgentStepsWithContext(
885829
new MockWebSocket() as unknown as WebSocket,
886830
{
887831
userInputId: 'test-user-input',
@@ -947,15 +891,10 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
947891
return 'mock-message-id'
948892
})
949893

950-
const mockCheckLiveUserInput = require('@codebuff/backend/live-user-inputs')
951-
spyOn(mockCheckLiveUserInput, 'checkLiveUserInput').mockImplementation(
952-
() => true,
953-
)
954-
955894
mockAgentState.output = undefined
956895
capturedAgentState = mockAgentState
957896

958-
const result = await loopAgentSteps(
897+
const result = await runLoopAgentStepsWithContext(
959898
new MockWebSocket() as unknown as WebSocket,
960899
{
961900
userInputId: 'test-user-input',

0 commit comments

Comments
 (0)