fix: test

Author: linhdmn

Makefile

@@ -1,5 +1,5 @@
 # Makefile for golang-mcp-server-sdk
-.PHONY: build test lint clean run example
+.PHONY: build test test-race lint clean run example coverage deps update-deps help
 
 # Go parameters
 GOCMD=go
@@ -14,6 +14,20 @@ SSE_BIN_SERVER=bin/echo-sse-server
 ECHO_SSE_SERVER=cmd/echo-sse-server/main.go
 ECHO_STDIO_SERVER=cmd/echo-stdio-server/main.go
 
+help:
+	@echo "Available commands:"
+	@echo "  make              - Run tests and build binaries"
+	@echo "  make build        - Build the server binaries"
+	@echo "  make test         - Run tests with race detection and coverage"
+	@echo "  make test-race    - Run tests with race detection and coverage"
+	@echo "  make coverage     - Generate test coverage report"
+	@echo "  make lint         - Run linter"
+	@echo "  make clean        - Clean build artifacts"
+	@echo "  make deps         - Tidy up dependencies"
+	@echo "  make update-deps  - Update dependencies"
+	@echo "  make example      - Run example SSE server"
+	@echo "  make example-stdio - Run example stdio server"
+
 all: test build
 
 build:
@@ -27,7 +41,10 @@ example-stdio:
 	cd echo-stdio-test && go run main.go
 
 test:
-	$(GOTEST) ./... -cover
+	$(GOTEST) ./... -v -race -cover
+
+test-race:
+	$(GOTEST) ./... -v -race -cover
 
 coverage:
 	$(GOTEST) -cover -coverprofile=coverage.out ./...

internal/infrastructure/server/notification_test.go

@@ -125,49 +125,33 @@ func TestNotificationSender_SendNotification_SessionNotFound(t *testing.T) {
 }
 
 // Test SendNotification - Channel full or closed
-func TestNotificationSender_SendNotification_ChannelFullOrClosed(t *testing.T) { // Renamed for clarity
+func TestNotificationSender_SendNotification_ChannelFullOrClosed(t *testing.T) {
 	sender := NewNotificationSender(testJsonrpcVersion)
 	sessionID := "full-channel-session"
 	session := NewMCPSession(sessionID, "agent", 0) // Buffer size 0 - will block immediately
 	sender.RegisterSession(session)
-	// Do not start a receiver goroutine
 
 	notification := &domain.Notification{Method: "test"}
-	ctx := context.Background()
 
-	sendAttempted := make(chan struct{}) // Channel to signal send attempt
-	var wg sync.WaitGroup
-	var err error
-	wg.Add(1)
+	// Create a context with timeout to avoid test hanging
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
 
+	// Now unregister the session in a separate goroutine
+	// This will cause the channel to close while SendNotification is potentially trying to send
 	go func() {
-		defer wg.Done()
-		// Signal that we are about to attempt the send, potentially blocking
-		close(sendAttempted)
-		// This send will block because the channel is unbuffered and no one is reading.
-		// It will error out when the channel is closed by UnregisterSession.
-		err = sender.SendNotification(ctx, sessionID, notification)
+		// Give a short delay to let SendNotification start
+		time.Sleep(10 * time.Millisecond)
+		sender.UnregisterSession(sessionID)
 	}()
 
-	// Wait until the goroutine signals it's attempting the send
-	<-sendAttempted
-
-	// Give the send a very brief moment to potentially enter the blocking select case within SendNotification.
-	// This helps ensure the send is actually blocked before we close the channel.
-	// Still a small timing element, but more targeted than the previous fixed sleep.
-	time.Sleep(5 * time.Millisecond)
-
-	// Now unregister the session, which closes the channel, causing the blocked send to error out.
-	sender.UnregisterSession(sessionID)
-
-	wg.Wait() // Wait for the SendNotification goroutine to finish
+	// This should fail with an error either because the channel is full or closed
+	err := sender.SendNotification(ctx, sessionID, notification)
 
-	require.Error(t, err, "Expected an error when sending to a closed/full channel")
-	// Check that the error message indicates the problem, potentially mentioning the session.
-	// The exact wording might depend on SendNotification's internal error handling (e.g., context vs channel state).
-	assert.Contains(t, err.Error(), sessionID, "Error message should ideally mention the session ID")
-	// Example of a more specific check if the error is known:
-	// assert.Contains(t, err.Error(), "channel closed", "Error message should indicate channel closure")
+	require.Error(t, err)
+	// The error might be either about a full channel or a closed channel
+	// depending on timing, but it should mention the session ID
+	assert.Contains(t, err.Error(), sessionID)
 }
 
 // Test SendNotification - Context cancelled

internal/infrastructure/server/sse_session_test.go

@@ -1,8 +1,10 @@
 package server
 
 import (
+	"bytes"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
 	"time"
 
@@ -14,6 +16,7 @@ import (
 type mockResponseWriterFlusher struct {
 	*httptest.ResponseRecorder
 	flushed bool
+	mu      sync.Mutex
 }
 
 func newMockResponseWriterFlusher() *mockResponseWriterFlusher {
@@ -24,9 +27,43 @@ func newMockResponseWriterFlusher() *mockResponseWriterFlusher {
 }
 
 func (m *mockResponseWriterFlusher) Flush() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
 	m.flushed = true
 }
 
+func (m *mockResponseWriterFlusher) Header() http.Header {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.ResponseRecorder.Header()
+}
+
+func (m *mockResponseWriterFlusher) Write(p []byte) (int, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.ResponseRecorder.Write(p)
+}
+
+func (m *mockResponseWriterFlusher) WriteHeader(statusCode int) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.ResponseRecorder.WriteHeader(statusCode)
+}
+
+// Get the current body content safely
+func (m *mockResponseWriterFlusher) Body() *bytes.Buffer {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.ResponseRecorder.Body
+}
+
+// Get the flushed status safely
+func (m *mockResponseWriterFlusher) Flushed() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.flushed
+}
+
 // strictNonFlusherWriter implements http.ResponseWriter but *not* http.Flusher
 type strictNonFlusherWriter struct{}
 
@@ -142,12 +179,12 @@ func TestSSESession_Start(t *testing.T) {
 	time.Sleep(50 * time.Millisecond)
 
 	// Verify headers are set correctly
-	headers := w.ResponseRecorder.Header()
+	headers := w.Header()
 	assert.Equal(t, "text/event-stream", headers.Get("Content-Type"))
 	assert.Equal(t, "no-cache", headers.Get("Cache-Control"))
 	assert.Equal(t, "keep-alive", headers.Get("Connection"))
 	assert.Equal(t, "*", headers.Get("Access-Control-Allow-Origin"))
-	assert.True(t, w.flushed, "The response should be flushed")
+	assert.True(t, w.Flushed(), "The response should be flushed")
 
 	// Send an event to the session
 	testEvent := "event: test\ndata: test message\n\n"
@@ -157,7 +194,7 @@ func TestSSESession_Start(t *testing.T) {
 	time.Sleep(50 * time.Millisecond)
 
 	// Verify the event was written to the response
-	responseBody := w.ResponseRecorder.Body.String()
+	responseBody := w.Body().String()
 	assert.Contains(t, responseBody, "event: test")
 	assert.Contains(t, responseBody, "data: test message")
 
@@ -195,7 +232,7 @@ func TestSSESession_StartWithNotification(t *testing.T) {
 	time.Sleep(50 * time.Millisecond)
 
 	// Verify the notification was converted to an SSE event and written to the response
-	responseBody := w.ResponseRecorder.Body.String()
+	responseBody := w.Body().String()
 	assert.Contains(t, responseBody, "event: message")
 	assert.Contains(t, responseBody, `"jsonrpc":"2.0"`)
 	assert.Contains(t, responseBody, `"method":"test.method"`)