first commit

internal/handlers/websocket/broadcast_test.go

@@ -0,0 +1,251 @@
+package websocket
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+// MockWebSocketHandler is a mock implementation for testing
+type MockWebSocketHandler struct {
+	mu         sync.Mutex
+	messages   []map[string]interface{}
+	broadcasts []string
+}
+
+func (m *MockWebSocketHandler) BroadcastMessage(messageType string, data interface{}) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.broadcasts = append(m.broadcasts, messageType)
+	m.messages = append(m.messages, map[string]interface{}{
+		"type": messageType,
+		"data": data,
+	})
+}
+
+func (m *MockWebSocketHandler) GetMessages() []map[string]interface{} {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	result := make([]map[string]interface{}, len(m.messages))
+	copy(result, m.messages)
+	return result
+}
+
+func (m *MockWebSocketHandler) GetBroadcasts() []string {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	result := make([]string, len(m.broadcasts))
+	copy(result, m.broadcasts)
+	return result
+}
+
+func (m *MockWebSocketHandler) Clear() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.messages = make([]map[string]interface{}, 0)
+	m.broadcasts = make([]string, 0)
+}
+
+func NewMockWebSocketHandler() *MockWebSocketHandler {
+	return &MockWebSocketHandler{
+		messages:   make([]map[string]interface{}, 0),
+		broadcasts: make([]string, 0),
+	}
+}
+
+func TestBroadcaster_StartHeartbeat(t *testing.T) {
+	mockHandler := NewMockWebSocketHandler()
+	broadcaster := NewBroadcaster(mockHandler)
+
+	// Start heartbeat with short interval for testing
+	broadcaster.StartHeartbeat(100 * time.Millisecond)
+
+	// Wait for a few heartbeats
+	time.Sleep(350 * time.Millisecond)
+
+	// Stop the broadcaster
+	broadcaster.Stop()
+
+	// Check if heartbeats were sent
+	messages := mockHandler.GetMessages()
+	if len(messages) == 0 {
+		t.Error("Expected heartbeat messages, but got none")
+	}
+
+	// Check that all messages are heartbeat type
+	broadcasts := mockHandler.GetBroadcasts()
+	for _, msgType := range broadcasts {
+		if msgType != "heartbeat" {
+			t.Errorf("Expected heartbeat message type, got %s", msgType)
+		}
+	}
+
+	t.Logf("Received %d heartbeat messages", len(messages))
+}
+
+func TestBroadcaster_BroadcastNotification(t *testing.T) {
+	mockHandler := NewMockWebSocketHandler()
+	broadcaster := NewBroadcaster(mockHandler)
+
+	// Send a notification
+	broadcaster.BroadcastNotification("Test Title", "Test Message", "info")
+
+	// Check if notification was sent
+	messages := mockHandler.GetMessages()
+	if len(messages) != 1 {
+		t.Errorf("Expected 1 message, got %d", len(messages))
+		return
+	}
+
+	msg := messages[0]
+	if msg["type"] != "notification" {
+		t.Errorf("Expected message type 'notification', got %s", msg["type"])
+	}
+
+	data := msg["data"].(map[string]interface{})
+	if data["title"] != "Test Title" {
+		t.Errorf("Expected title 'Test Title', got %s", data["title"])
+	}
+	if data["message"] != "Test Message" {
+		t.Errorf("Expected message 'Test Message', got %s", data["message"])
+	}
+	if data["level"] != "info" {
+		t.Errorf("Expected level 'info', got %s", data["level"])
+	}
+
+	t.Logf("Notification sent successfully: %+v", data)
+}
+
+func TestBroadcaster_SimulateDataStream(t *testing.T) {
+	mockHandler := NewMockWebSocketHandler()
+	broadcaster := NewBroadcaster(mockHandler)
+
+	// Start data stream with short interval for testing
+	broadcaster.SimulateDataStream()
+
+	// Wait for a few data points
+	time.Sleep(550 * time.Millisecond)
+
+	// Stop the broadcaster
+	broadcaster.Stop()
+
+	// Check if data stream messages were sent
+	messages := mockHandler.GetMessages()
+	if len(messages) == 0 {
+		t.Error("Expected data stream messages, but got none")
+	}
+
+	// Check that all messages are data_stream type
+	broadcasts := mockHandler.GetBroadcasts()
+	for _, msgType := range broadcasts {
+		if msgType != "data_stream" {
+			t.Errorf("Expected data_stream message type, got %s", msgType)
+		}
+	}
+
+	// Check data structure
+	for i, msg := range messages {
+		data := msg["data"].(map[string]interface{})
+		if data["type"] != "simulated_data" {
+			t.Errorf("Expected data type 'simulated_data', got %s", data["type"])
+		}
+		if id, ok := data["id"].(int); ok {
+			if id != i+1 {
+				t.Errorf("Expected id %d, got %d", i+1, id)
+			}
+		}
+		if value, ok := data["value"].(int); ok {
+			expectedValue := (i + 1) * 10
+			if value != expectedValue {
+				t.Errorf("Expected value %d, got %d", expectedValue, value)
+			}
+		}
+	}
+
+	t.Logf("Received %d data stream messages", len(messages))
+}
+
+func TestBroadcaster_Stop(t *testing.T) {
+	mockHandler := NewMockWebSocketHandler()
+	broadcaster := NewBroadcaster(mockHandler)
+
+	// Start heartbeat
+	broadcaster.StartHeartbeat(50 * time.Millisecond)
+
+	// Wait a bit
+	time.Sleep(100 * time.Millisecond)
+
+	// Stop the broadcaster
+	broadcaster.Stop()
+
+	// Clear previous messages
+	mockHandler.Clear()
+
+	// Wait a bit more to ensure no new messages are sent
+	time.Sleep(200 * time.Millisecond)
+
+	// Check that no new messages were sent after stopping
+	messages := mockHandler.GetMessages()
+	if len(messages) > 0 {
+		t.Errorf("Expected no messages after stopping, but got %d", len(messages))
+	}
+
+	// Clear quit channel to allow reuse in tests
+	broadcaster.quit = make(chan struct{})
+
+	t.Log("Broadcaster stopped successfully")
+}
+
+func TestBroadcaster_MultipleOperations(t *testing.T) {
+	mockHandler := NewMockWebSocketHandler()
+	broadcaster := NewBroadcaster(mockHandler)
+
+	// Start heartbeat
+	broadcaster.StartHeartbeat(100 * time.Millisecond)
+
+	// Send notification
+	broadcaster.BroadcastNotification("Test", "Message", "warning")
+
+	// Start data stream
+	broadcaster.SimulateDataStream()
+
+	// Wait for some activity
+	time.Sleep(350 * time.Millisecond)
+
+	// Stop everything
+	broadcaster.Stop()
+
+	// Check results
+	messages := mockHandler.GetMessages()
+	if len(messages) == 0 {
+		t.Error("Expected messages from multiple operations, but got none")
+	}
+
+	broadcasts := mockHandler.GetBroadcasts()
+	hasHeartbeat := false
+	hasNotification := false
+	hasDataStream := false
+
+	for _, msgType := range broadcasts {
+		switch msgType {
+		case "heartbeat":
+			hasHeartbeat = true
+		case "notification":
+			hasNotification = true
+		case "data_stream":
+			hasDataStream = true
+		}
+	}
+
+	if !hasHeartbeat {
+		t.Error("Expected heartbeat messages")
+	}
+	if !hasNotification {
+		t.Error("Expected notification message")
+	}
+	if !hasDataStream {
+		t.Error("Expected data stream messages")
+	}
+
+	t.Logf("Multiple operations test passed: %d total messages", len(messages))
+}