at-container-registry/pkg/appview/handlers/device_test.go

package handlers

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

	"atcr.io/pkg/appview/db"
	"github.com/go-chi/chi/v5"
	_ "github.com/mattn/go-sqlite3"
)

// setupTestDB creates an in-memory SQLite database with full schema for testing
func setupTestDB(t *testing.T) *sql.DB {
	database, err := db.InitDB(":memory:", true)
	if err != nil {
		t.Fatalf("Failed to initialize test database: %v", err)
	}
	return database
}

// Test getClientIP function (existing test, expanded)
func TestGetClientIP(t *testing.T) {
	tests := []struct {
		name          string
		remoteAddr    string
		xForwardedFor string
		xRealIP       string
		expectedIP    string
	}{
		{
			name:          "X-Forwarded-For single IP",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "10.0.0.1",
			xRealIP:       "",
			expectedIP:    "10.0.0.1",
		},
		{
			name:          "X-Forwarded-For multiple IPs",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "10.0.0.1, 10.0.0.2, 10.0.0.3",
			xRealIP:       "",
			expectedIP:    "10.0.0.1",
		},
		{
			name:          "X-Forwarded-For with whitespace",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "  10.0.0.1  ",
			xRealIP:       "",
			expectedIP:    "10.0.0.1",
		},
		{
			name:          "X-Real-IP when no X-Forwarded-For",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "",
			xRealIP:       "10.0.0.2",
			expectedIP:    "10.0.0.2",
		},
		{
			name:          "X-Forwarded-For takes priority over X-Real-IP",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "10.0.0.1",
			xRealIP:       "10.0.0.2",
			expectedIP:    "10.0.0.1",
		},
		{
			name:          "RemoteAddr fallback with port",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "",
			xRealIP:       "",
			expectedIP:    "192.168.1.1",
		},
		{
			name:          "RemoteAddr fallback without port",
			remoteAddr:    "192.168.1.1",
			xForwardedFor: "",
			xRealIP:       "",
			expectedIP:    "192.168.1.1",
		},
		{
			name:          "IPv6 RemoteAddr",
			remoteAddr:    "[::1]:1234",
			xForwardedFor: "",
			xRealIP:       "",
			expectedIP:    "[",
		},
		{
			name:          "IPv6 in X-Forwarded-For",
			remoteAddr:    "192.168.1.1:1234",
			xForwardedFor: "2001:db8::1",
			xRealIP:       "",
			expectedIP:    "2001:db8::1",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest("GET", "http://example.com/test", nil)
			req.RemoteAddr = tt.remoteAddr

			if tt.xForwardedFor != "" {
				req.Header.Set("X-Forwarded-For", tt.xForwardedFor)
			}

			if tt.xRealIP != "" {
				req.Header.Set("X-Real-IP", tt.xRealIP)
			}

			result := getClientIP(req)
			if result != tt.expectedIP {
				t.Errorf("getClientIP() = %q, want %q", result, tt.expectedIP)
			}
		})
	}
}

func TestDeviceCodeHandler_Success(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceCodeHandler{
		Store:          store,
		AppViewBaseURL: "http://localhost:5000",
	}

	reqBody := DeviceCodeRequest{
		DeviceName: "My Test Device",
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/auth/device/code", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var response DeviceCodeResponse
	if err := json.NewDecoder(rr.Body).Decode(&response); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if response.DeviceCode == "" {
		t.Error("Expected device_code to be set")
	}
	if response.UserCode == "" {
		t.Error("Expected user_code to be set")
	}
	if !strings.HasPrefix(response.VerificationURI, "http://localhost:5000") {
		t.Errorf("Expected verification_uri to start with base URL, got %s", response.VerificationURI)
	}
	if response.ExpiresIn != 600 {
		t.Errorf("Expected expires_in to be 600, got %d", response.ExpiresIn)
	}
	if response.Interval != 5 {
		t.Errorf("Expected interval to be 5, got %d", response.Interval)
	}
}

func TestDeviceCodeHandler_DefaultDeviceName(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceCodeHandler{
		Store:          store,
		AppViewBaseURL: "http://localhost:5000",
	}

	// Empty device name should get default
	reqBody := DeviceCodeRequest{
		DeviceName: "",
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/auth/device/code", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var response DeviceCodeResponse
	if err := json.NewDecoder(rr.Body).Decode(&response); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if response.UserCode == "" {
		t.Error("Expected user_code to be set even with default device name")
	}
}

func TestDeviceCodeHandler_MethodNotAllowed(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceCodeHandler{
		Store:          store,
		AppViewBaseURL: "http://localhost:5000",
	}

	req := httptest.NewRequest("GET", "/auth/device/code", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
	}
}

func TestDeviceTokenHandler_AuthorizationPending(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceTokenHandler{
		Store: store,
	}

	// Create a pending authorization
	pending, err := store.CreatePendingAuth("Test Device", "127.0.0.1", "TestAgent/1.0")
	if err != nil {
		t.Fatalf("Failed to create pending auth: %v", err)
	}

	// Poll before approval
	reqBody := DeviceTokenRequest{
		DeviceCode: pending.DeviceCode,
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/auth/device/token", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var response DeviceTokenResponse
	if err := json.NewDecoder(rr.Body).Decode(&response); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if response.Error != "authorization_pending" {
		t.Errorf("Expected error 'authorization_pending', got %s", response.Error)
	}
}

func TestDeviceTokenHandler_ExpiredToken(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceTokenHandler{
		Store: store,
	}

	// Try to poll with invalid device code
	reqBody := DeviceTokenRequest{
		DeviceCode: "invalid_code_12345",
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/auth/device/token", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var response DeviceTokenResponse
	if err := json.NewDecoder(rr.Body).Decode(&response); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if response.Error != "expired_token" {
		t.Errorf("Expected error 'expired_token', got %s", response.Error)
	}
}

func TestDeviceTokenHandler_Approved(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceTokenHandler{
		Store: store,
	}

	// Create a pending authorization
	pending, err := store.CreatePendingAuth("Test Device", "127.0.0.1", "TestAgent/1.0")
	if err != nil {
		t.Fatalf("Failed to create pending auth: %v", err)
	}

	// Create user first (required for foreign key)
	_, err = database.Exec(`
		INSERT INTO users (did, handle, pds_endpoint, last_seen)
		VALUES (?, ?, ?, ?)
	`, "did:plc:test123", "test.bsky.social", "https://bsky.social", time.Now())
	if err != nil {
		t.Fatalf("Failed to create user: %v", err)
	}

	// Approve it
	_, err = store.ApprovePending(pending.UserCode, "did:plc:test123", "test.bsky.social")
	if err != nil {
		t.Fatalf("Failed to approve pending: %v", err)
	}

	// Poll after approval
	reqBody := DeviceTokenRequest{
		DeviceCode: pending.DeviceCode,
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/auth/device/token", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var response DeviceTokenResponse
	if err := json.NewDecoder(rr.Body).Decode(&response); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if response.Error != "" {
		t.Errorf("Expected no error, got %s", response.Error)
	}
	if response.DeviceSecret == "" {
		t.Error("Expected device_secret to be set")
	}
	if response.DID != "did:plc:test123" {
		t.Errorf("Expected DID 'did:plc:test123', got %s", response.DID)
	}
}

func TestDeviceTokenHandler_MethodNotAllowed(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	handler := &DeviceTokenHandler{
		Store: store,
	}

	req := httptest.NewRequest("GET", "/auth/device/token", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
	}
}

func TestDeviceApprovalPageHandler_NotLoggedIn(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &DeviceApprovalPageHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("GET", "/device?user_code=ABC123", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	// Should redirect to login
	if rr.Code != http.StatusFound {
		t.Errorf("Expected status %d, got %d", http.StatusFound, rr.Code)
	}

	location := rr.Header().Get("Location")
	if !strings.Contains(location, "/auth/oauth/login") {
		t.Errorf("Expected redirect to login, got %s", location)
	}
}

func TestDeviceApprovalPageHandler_MissingUserCode(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	// Create user first (required for foreign key)
	_, err := database.Exec(`
		INSERT INTO users (did, handle, pds_endpoint, last_seen)
		VALUES (?, ?, ?, ?)
	`, "did:plc:test123", "test.bsky.social", "https://bsky.social", time.Now())
	if err != nil {
		t.Fatalf("Failed to create user: %v", err)
	}

	// Create a session
	sessionID, _ := sessionStore.Create("did:plc:test123", "test.bsky.social", "https://pds.example.com", 24*time.Hour)

	handler := &DeviceApprovalPageHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("GET", "/device", nil) // No user_code parameter
	req.AddCookie(&http.Cookie{
		Name:  "atcr_session",
		Value: sessionID,
	})

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusBadRequest {
		t.Errorf("Expected status %d, got %d", http.StatusBadRequest, rr.Code)
	}
}

func TestDeviceApprovalPageHandler_MethodNotAllowed(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &DeviceApprovalPageHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("POST", "/device?user_code=ABC123", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
	}
}

func TestDeviceApproveHandler_Unauthorized(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &DeviceApproveHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	reqBody := DeviceApproveRequest{
		UserCode: "ABC123",
		Approve:  true,
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/device/approve", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Errorf("Expected status %d, got %d", http.StatusUnauthorized, rr.Code)
	}
}

func TestDeviceApproveHandler_Deny(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	// Create user first (required for foreign key)
	_, err := database.Exec(`
		INSERT INTO users (did, handle, pds_endpoint, last_seen)
		VALUES (?, ?, ?, ?)
	`, "did:plc:test123", "test.bsky.social", "https://bsky.social", time.Now())
	if err != nil {
		t.Fatalf("Failed to create user: %v", err)
	}

	// Create a session
	sessionID, _ := sessionStore.Create("did:plc:test123", "test.bsky.social", "https://pds.example.com", 24*time.Hour)

	handler := &DeviceApproveHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	reqBody := DeviceApproveRequest{
		UserCode: "ABC123",
		Approve:  false,
	}
	body, _ := json.Marshal(reqBody)
	req := httptest.NewRequest("POST", "/device/approve", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.AddCookie(&http.Cookie{
		Name:  "atcr_session",
		Value: sessionID,
	})

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var response map[string]string
	if err := json.NewDecoder(rr.Body).Decode(&response); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if response["status"] != "denied" {
		t.Errorf("Expected status 'denied', got %s", response["status"])
	}
}

func TestDeviceApproveHandler_MethodNotAllowed(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &DeviceApproveHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("GET", "/device/approve", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
	}
}

func TestListDevicesHandler_Unauthorized(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &ListDevicesHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("GET", "/api/devices", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Errorf("Expected status %d, got %d", http.StatusUnauthorized, rr.Code)
	}
}

func TestListDevicesHandler_Success(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	// Create user first (required for foreign key)
	_, err := database.Exec(`
		INSERT INTO users (did, handle, pds_endpoint, last_seen)
		VALUES (?, ?, ?, ?)
	`, "did:plc:test123", "test.bsky.social", "https://bsky.social", time.Now())
	if err != nil {
		t.Fatalf("Failed to create user: %v", err)
	}

	// Create a session
	sessionID, _ := sessionStore.Create("did:plc:test123", "test.bsky.social", "https://pds.example.com", 24*time.Hour)

	// Create some devices
	pending, _ := store.CreatePendingAuth("Device 1", "127.0.0.1", "TestAgent/1.0")
	store.ApprovePending(pending.UserCode, "did:plc:test123", "test.bsky.social")

	handler := &ListDevicesHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("GET", "/api/devices", nil)
	req.AddCookie(&http.Cookie{
		Name:  "atcr_session",
		Value: sessionID,
	})

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("Expected status %d, got %d", http.StatusOK, rr.Code)
	}

	var devices []db.Device
	if err := json.NewDecoder(rr.Body).Decode(&devices); err != nil {
		t.Fatalf("Failed to decode response: %v", err)
	}

	if len(devices) != 1 {
		t.Errorf("Expected 1 device, got %d", len(devices))
	}
}

func TestListDevicesHandler_MethodNotAllowed(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &ListDevicesHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("POST", "/api/devices", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
	}
}

func TestRevokeDeviceHandler_Unauthorized(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &RevokeDeviceHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("DELETE", "/api/devices/device123", nil)

	// Add chi URL parameter
	rctx := chi.NewRouteContext()
	rctx.URLParams.Add("id", "device123")
	req = req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx))

	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Errorf("Expected status %d, got %d", http.StatusUnauthorized, rr.Code)
	}
}

func TestRevokeDeviceHandler_MethodNotAllowed(t *testing.T) {
	database := setupTestDB(t)
	defer database.Close()

	store := db.NewDeviceStore(database)
	sessionStore := db.NewSessionStore(database)

	handler := &RevokeDeviceHandler{
		Store:        store,
		SessionStore: sessionStore,
	}

	req := httptest.NewRequest("GET", "/api/devices/device123", nil)
	rr := httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Errorf("Expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
	}
}