chore(push-gateway): ai02 module updates - config, hub, ws, kafka, health, docs

services/push-gateway/internal/ws/handler_test.go

@@ -0,0 +1,196 @@
+package ws
+
+import (
+	"encoding/json"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/edu-cloud/shared-go/jwks"
+	"github.com/golang-jwt/jwt/v5"
+)
+
+// TestBuildMessage verifies the WebSocket message envelope format (02 §4.1).
+func TestBuildMessage(t *testing.T) {
+	data := map[string]any{"title": "hello", "body": "world"}
+	msg, err := buildMessage("notification.created", data)
+	if err != nil {
+		t.Fatalf("buildMessage: %v", err)
+	}
+
+	var payload map[string]any
+	if err := json.Unmarshal(msg, &payload); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	if payload["type"] != "message" {
+		t.Errorf("type = %v, want message", payload["type"])
+	}
+	if payload["event"] != "notification.created" {
+		t.Errorf("event = %v, want notification.created", payload["event"])
+	}
+	if payload["timestamp"] == nil {
+		t.Error("timestamp missing")
+	}
+	dataField, ok := payload["data"].(map[string]any)
+	if !ok {
+		t.Fatalf("data should be a map, got %T", payload["data"])
+	}
+	if dataField["title"] != "hello" {
+		t.Errorf("data.title = %v, want hello", dataField["title"])
+	}
+}
+
+// TestBuildMessageNilData verifies nil data is handled gracefully.
+func TestBuildMessageNilData(t *testing.T) {
+	msg, err := buildMessage("test.event", nil)
+	if err != nil {
+		t.Fatalf("buildMessage: %v", err)
+	}
+	var payload map[string]any
+	if err := json.Unmarshal(msg, &payload); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	if payload["data"] != nil {
+		t.Errorf("data = %v, want nil", payload["data"])
+	}
+}
+
+// TestMustMarshal verifies best-effort JSON encoding.
+func TestMustMarshal(t *testing.T) {
+	// nil -> "null"
+	if got := mustMarshal(nil); string(got) != "null" {
+		t.Errorf("mustMarshal(nil) = %q, want null", got)
+	}
+	// valid map -> JSON
+	got := mustMarshal(map[string]any{"a": 1})
+	var m map[string]any
+	if err := json.Unmarshal(got, &m); err != nil {
+		t.Fatalf("mustMarshal result unmarshal: %v", err)
+	}
+	if m["a"].(float64) != 1 {
+		t.Errorf("mustMarshal.a = %v, want 1", m["a"])
+	}
+}
+
+// TestBuildOriginChecker verifies origin whitelist logic.
+func TestBuildOriginChecker(t *testing.T) {
+	allowed := map[string]struct{}{
+		"http://localhost:3000": {},
+		"https://app.example.com": {},
+	}
+
+	// Non-devMode: strict whitelist.
+	checker := buildOriginChecker(allowed, false)
+
+	tests := []struct {
+		name   string
+		origin string
+		want   bool
+	}{
+		{"empty origin (non-browser)", "", true},
+		{"allowed origin", "http://localhost:3000", true},
+		{"another allowed", "https://app.example.com", true},
+		{"disallowed origin", "http://evil.com", false},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			req := httptest.NewRequest("GET", "/ws", nil)
+			if tc.origin != "" {
+				req.Header.Set("Origin", tc.origin)
+			}
+			if got := checker(req); got != tc.want {
+				t.Errorf("checker(origin=%q) = %v, want %v", tc.origin, got, tc.want)
+			}
+		})
+	}
+
+	// DevMode with empty whitelist: allow all.
+	devChecker := buildOriginChecker(map[string]struct{}{}, true)
+	req := httptest.NewRequest("GET", "/ws", nil)
+	req.Header.Set("Origin", "http://anything.com")
+	if !devChecker(req) {
+		t.Error("devMode empty whitelist should allow all origins")
+	}
+}
+
+// TestValidateHS256 verifies HS256 JWT validation in DevMode fallback.
+func TestValidateHS256(t *testing.T) {
+	secret := "test-secret"
+
+	// Valid token with user_id claim.
+	claims := jwks.Claims{
+		UserID: "user-123",
+		RegisteredClaims: jwt.RegisteredClaims{
+			ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
+		},
+	}
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	tokenStr, err := token.SignedString([]byte(secret))
+	if err != nil {
+		t.Fatalf("SignedString: %v", err)
+	}
+
+	userID, err := validateHS256(tokenStr, secret)
+	if err != nil {
+		t.Fatalf("validateHS256: %v", err)
+	}
+	if userID != "user-123" {
+		t.Errorf("userID = %q, want user-123", userID)
+	}
+
+	// Wrong secret -> error.
+	_, err = validateHS256(tokenStr, "wrong-secret")
+	if err == nil {
+		t.Error("validateHS256 with wrong secret should fail")
+	}
+
+	// Expired token -> error.
+	expiredClaims := jwks.Claims{
+		UserID: "user-exp",
+		RegisteredClaims: jwt.RegisteredClaims{
+			ExpiresAt: jwt.NewNumericDate(time.Now().Add(-1 * time.Hour)),
+		},
+	}
+	expiredToken := jwt.NewWithClaims(jwt.SigningMethodHS256, expiredClaims)
+	expiredStr, _ := expiredToken.SignedString([]byte(secret))
+	_, err = validateHS256(expiredStr, secret)
+	if err == nil {
+		t.Error("validateHS256 with expired token should fail")
+	}
+
+	// Missing user_id -> error.
+	noUserClaims := jwks.Claims{
+		RegisteredClaims: jwt.RegisteredClaims{
+			ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
+		},
+	}
+	noUserToken := jwt.NewWithClaims(jwt.SigningMethodHS256, noUserClaims)
+	noUserStr, _ := noUserToken.SignedString([]byte(secret))
+	_, err = validateHS256(noUserStr, secret)
+	if err != errMissingUserID {
+		t.Errorf("validateHS256 without user_id err = %v, want errMissingUserID", err)
+	}
+}
+
+// TestContextWithTimeout verifies the helper returns a working context.
+func TestContextWithTimeout(t *testing.T) {
+	ctx, cancel := contextWithTimeout(100 * time.Millisecond)
+	defer cancel()
+	if ctx == nil {
+		t.Fatal("context is nil")
+	}
+	// Context should not be done immediately.
+	select {
+	case <-ctx.Done():
+		t.Fatal("context done immediately")
+	default:
+	}
+	// Wait for timeout.
+	time.Sleep(150 * time.Millisecond)
+	select {
+	case <-ctx.Done():
+		// expected
+	default:
+		t.Error("context not done after timeout")
+	}
+}