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

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