From 00b6b5ee0ea9f2ba19cf1dcd97508c753a5ff27c Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 26 Jan 2026 18:40:06 +0800 Subject: [PATCH] Enhance Tests with SafeWriter Closure for Async Writes - Added calls to CloseSafeWriter in multiple test cases to ensure all asynchronous writes are completed before verifying output. - Improved test reliability by ensuring that the SafeWriter's buffer is properly flushed, preventing potential race conditions in output verification. --- agent/context/jsapi_output_test.go | 15 +++++++++++++++ agent/context/message_events_test.go | 4 ++++ 2 files changed, 19 insertions(+) diff --git a/agent/context/jsapi_output_test.go b/agent/context/jsapi_output_test.go index 128b18a4..91592ffd 100644 --- a/agent/context/jsapi_output_test.go +++ b/agent/context/jsapi_output_test.go @@ -771,6 +771,9 @@ func TestJsValueEndBlock(t *testing.T) { } assert.Equal(t, true, result["success"], "EndBlock should work correctly") + // Close SafeWriter to wait for all async writes to complete + cxt.CloseSafeWriter() + // Verify that block_end event was sent output := mockWriter.buffer.String() assert.Contains(t, output, "block_end", "Output should contain block_end event") @@ -825,6 +828,9 @@ func TestJsValueSendStream(t *testing.T) { } assert.Equal(t, true, result["success"], "SendStream should work correctly") + // Close SafeWriter to wait for all async writes to complete + cxt.CloseSafeWriter() + // Verify message_start was sent but NOT message_end output := mockWriter.buffer.String() assert.Contains(t, output, "message_start", "Output should contain message_start event") @@ -873,6 +879,9 @@ func TestJsValueSendStreamWithBlockID(t *testing.T) { } assert.Equal(t, true, result["success"], "SendStream with blockId should succeed") + // Close SafeWriter to wait for all async writes to complete + cxt.CloseSafeWriter() + // Verify block_start was also sent output := mockWriter.buffer.String() assert.Contains(t, output, "block_start", "Output should contain block_start event") @@ -922,6 +931,9 @@ func TestJsValueEnd(t *testing.T) { } assert.Equal(t, true, result["success"], "End should work correctly") + // Close SafeWriter to wait for all async writes to complete + cxt.CloseSafeWriter() + // Verify message_end was sent output := mockWriter.buffer.String() assert.Contains(t, output, "message_end", "Output should contain message_end event after End()") @@ -971,6 +983,9 @@ func TestJsValueEndWithFinalContent(t *testing.T) { } assert.Equal(t, true, result["success"], "End with final content should work correctly") + // Close SafeWriter to wait for all async writes to complete + cxt.CloseSafeWriter() + // Verify message_end was sent output := mockWriter.buffer.String() assert.Contains(t, output, "message_end", "Output should contain message_end event") diff --git a/agent/context/message_events_test.go b/agent/context/message_events_test.go index b06ae1f6..efb9ace5 100644 --- a/agent/context/message_events_test.go +++ b/agent/context/message_events_test.go @@ -39,6 +39,10 @@ func TestMessageLifecycleEvents(t *testing.T) { // Flush to ensure all messages are written ctx.Flush() + // Close SafeWriter to wait for all async writes to complete + // SafeWriter uses a channel-based queue, so we must close it before reading buffer + ctx.CloseSafeWriter() + // Parse output to find events output := buf.String() t.Logf("Output:\n%s", output)