baron-sso/backend/internal/handler/auth_handler_qr_test.go

package handler

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/stretchr/testify/assert"
)

// --- Mock Redis ---

type mockRedisRepo struct {
	data map[string]string
}

func (m *mockRedisRepo) Set(key, value string, ttl time.Duration) error {
	if m.data == nil {
		m.data = make(map[string]string)
	}
	m.data[key] = value
	return nil
}

func (m *mockRedisRepo) Get(key string) (string, error) {
	// Bypass rate limiting for tests
	if strings.HasPrefix(key, "poll_meta:") {
		return "", nil
	}
	return m.data[key], nil
}

func (m *mockRedisRepo) Delete(key string) error {
	delete(m.data, key)
	return nil
}

func (m *mockRedisRepo) StoreVerificationCode(phone, code string) error {
	return m.Set("sms:"+phone, code, time.Minute)
}

func (m *mockRedisRepo) GetVerificationCode(phone string) (string, error) {
	return m.Get("sms:" + phone)
}

func (m *mockRedisRepo) DeleteVerificationCode(phone string) error {
	return m.Delete("sms:" + phone)
}

// --- Tests ---

func TestQRLoginFlow_Success(t *testing.T) {
	redis := &mockRedisRepo{data: make(map[string]string)}
	h := &AuthHandler{
		RedisService: redis,
	}
	app := fiber.New()
	app.Post("/api/v1/auth/qr/init", h.InitQRLogin)
	app.Post("/api/v1/auth/qr/poll", h.PollQRLogin)

	// 1. Init QR Login
	req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/qr/init", nil)
	resp, _ := app.Test(req, -1)
	assert.Equal(t, http.StatusOK, resp.StatusCode)

	var initResp map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&initResp)
	pendingRef := initResp["pendingRef"].(string)

	// 2. Poll (Pending)
	body, _ := json.Marshal(map[string]string{"pendingRef": pendingRef})
	req = httptest.NewRequest(http.MethodPost, "/api/v1/auth/qr/poll", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, _ = app.Test(req, -1)

	// Expect authorization_pending (400)
	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
	var pollResp map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&pollResp)
	assert.Equal(t, "authorization_pending", pollResp["error"])
	assert.Equal(t, "authorization_pending", pollResp["code"])

	// 3. Mock Approval
	sessionData, _ := json.Marshal(map[string]string{
		"status": "success",
		"jwt":    "mock-session-jwt",
	})
	redis.data["enchanted_session:"+pendingRef] = string(sessionData)

	// 4. Poll (Success)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/auth/qr/poll", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, _ = app.Test(req, -1)

	assert.Equal(t, http.StatusOK, resp.StatusCode)

	var successResp map[string]interface{}
	json.NewDecoder(resp.Body).Decode(&successResp)
	assert.Equal(t, "ok", successResp["status"])
	assert.Equal(t, "mock-session-jwt", successResp["sessionJwt"])
}

func TestScanQRLogin_Success(t *testing.T) {
	redis := &mockRedisRepo{data: make(map[string]string)}
	idp := &mockIdpProvider{userExists: true}

	h := &AuthHandler{
		RedisService: redis,
		IdpProvider:  idp,
	}
	app := fiber.New()
	app.Post("/api/v1/auth/qr/approve", h.ScanQRLogin)

	pendingRef := "test-ref"
	redis.data["enchanted_session:"+pendingRef] = `{"status":"pending"}`

	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		if r.URL.Path == "/sessions/whoami" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"identity": map[string]interface{}{
					"id": "user-123",
					"traits": map[string]interface{}{
						"email": "user@example.com",
					},
				},
			}), nil
		}
		return httpResponse(r, http.StatusNotFound, "not found"), nil
	})

	origDefault := http.DefaultClient
	http.DefaultClient = &http.Client{Transport: transport}
	defer func() { http.DefaultClient = origDefault }()

	body, _ := json.Marshal(map[string]string{
		"pendingRef": pendingRef,
		"token":      "valid-token",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/qr/approve", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, _ := app.Test(req, -1)

	assert.Equal(t, http.StatusOK, resp.StatusCode)
}

func TestResolveConsentSubjects_TokenAndCookie(t *testing.T) {
	h := &AuthHandler{}

	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		if r.Header.Get("X-Session-Token") == "token-123" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"identity": map[string]interface{}{
					"id": "user-token",
					"traits": map[string]interface{}{
						"email": "token@test.com",
					},
				},
			}), nil
		}
		if r.Header.Get("Cookie") == "ory_kratos_session=cookie-123" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"identity": map[string]interface{}{
					"id": "user-cookie",
					"traits": map[string]interface{}{
						"email": "cookie@test.com",
						"phone": "010-1234-5678",
					},
				},
			}), nil
		}
		return httpResponse(r, http.StatusUnauthorized, "unauthorized"), nil
	})

	origDefault := http.DefaultClient
	http.DefaultClient = &http.Client{Transport: transport}
	defer func() { http.DefaultClient = origDefault }()

	app := fiber.New()

	// Token case
	app.Get("/test-token", func(c *fiber.Ctx) error {
		subjects, err := h.resolveConsentSubjects(c)
		assert.NoError(t, err)
		assert.Contains(t, subjects, "user-token")
		return c.SendStatus(200)
	})
	req := httptest.NewRequest("GET", "/test-token", nil)
	req.Header.Set("Authorization", "Bearer token-123")
	app.Test(req, -1)

	// Cookie case
	app.Get("/test-cookie", func(c *fiber.Ctx) error {
		subjects, err := h.resolveConsentSubjects(c)
		assert.NoError(t, err)
		assert.Contains(t, subjects, "user-cookie")
		return c.SendStatus(200)
	})
	req = httptest.NewRequest("GET", "/test-cookie", nil)
	req.Header.Set("Cookie", "ory_kratos_session=cookie-123")
	app.Test(req, -1)
}