first commit

internal/handlers/websocket/broadcast.go

@@ -0,0 +1,111 @@
+package websocket
+
+import (
+	"sync"
+	"time"
+)
+
+// WebSocketBroadcaster defines the interface for broadcasting messages
+type WebSocketBroadcaster interface {
+	BroadcastMessage(messageType string, data interface{})
+}
+
+// Broadcaster handles server-initiated broadcasts to WebSocket clients
+type Broadcaster struct {
+	handler WebSocketBroadcaster
+	tickers []*time.Ticker
+	quit    chan struct{}
+	mu      sync.Mutex
+}
+
+// NewBroadcaster creates a new Broadcaster instance
+func NewBroadcaster(handler WebSocketBroadcaster) *Broadcaster {
+	return &Broadcaster{
+		handler: handler,
+		tickers: make([]*time.Ticker, 0),
+		quit:    make(chan struct{}),
+	}
+}
+
+// StartHeartbeat starts sending periodic heartbeat messages to all clients
+func (b *Broadcaster) StartHeartbeat(interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	b.tickers = append(b.tickers, ticker)
+	go func() {
+		defer func() {
+			// Remove ticker from slice when done
+			for i, t := range b.tickers {
+				if t == ticker {
+					b.tickers = append(b.tickers[:i], b.tickers[i+1:]...)
+					break
+				}
+			}
+		}()
+		for {
+			select {
+			case <-ticker.C:
+				b.handler.BroadcastMessage("heartbeat", map[string]interface{}{
+					"message":   "Server heartbeat",
+					"timestamp": time.Now().Format(time.RFC3339),
+				})
+			case <-b.quit:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+}
+
+// Stop stops the broadcaster
+func (b *Broadcaster) Stop() {
+	close(b.quit)
+	for _, ticker := range b.tickers {
+		if ticker != nil {
+			ticker.Stop()
+		}
+	}
+	b.tickers = nil
+}
+
+// BroadcastNotification sends a notification message to all clients
+func (b *Broadcaster) BroadcastNotification(title, message, level string) {
+	b.handler.BroadcastMessage("notification", map[string]interface{}{
+		"title":   title,
+		"message": message,
+		"level":   level,
+		"time":    time.Now().Format(time.RFC3339),
+	})
+}
+
+// SimulateDataStream simulates streaming data to clients (useful for demos)
+func (b *Broadcaster) SimulateDataStream() {
+	ticker := time.NewTicker(100 * time.Millisecond)
+	b.tickers = append(b.tickers, ticker)
+	go func() {
+		defer func() {
+			// Remove ticker from slice when done
+			for i, t := range b.tickers {
+				if t == ticker {
+					b.tickers = append(b.tickers[:i], b.tickers[i+1:]...)
+					break
+				}
+			}
+		}()
+		counter := 0
+		for {
+			select {
+			case <-ticker.C:
+				counter++
+				b.handler.BroadcastMessage("data_stream", map[string]interface{}{
+					"id":        counter,
+					"value":     counter * 10,
+					"timestamp": time.Now().Format(time.RFC3339),
+					"type":      "simulated_data",
+				})
+			case <-b.quit:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+}

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))
+}