veza/veza-backend-api/internal/handlers/two_factor_handler_test.go

package handlers

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"veza-backend-api/internal/models"
	"veza-backend-api/internal/services"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"go.uber.org/zap"
)

// MockTwoFactorService mocks TwoFactorService
type MockTwoFactorService struct {
	mock.Mock
}

func (m *MockTwoFactorService) GetTwoFactorStatus(ctx context.Context, userID uuid.UUID) (bool, error) {
	args := m.Called(ctx, userID)
	return args.Bool(0), args.Error(1)
}

func (m *MockTwoFactorService) GenerateSecret(user *models.User) (*services.TwoFactorSetup, error) {
	args := m.Called(user)
	if args.Get(0) == nil {
		return nil, args.Error(1)
	}
	return args.Get(0).(*services.TwoFactorSetup), args.Error(1)
}

func (m *MockTwoFactorService) VerifyTOTPCode(secret, code string) bool {
	args := m.Called(secret, code)
	return args.Bool(0)
}

func (m *MockTwoFactorService) GenerateRecoveryCodes() []string {
	args := m.Called()
	return args.Get(0).([]string)
}

func (m *MockTwoFactorService) EnableTwoFactor(ctx context.Context, userID uuid.UUID, secret string, recoveryCodes []string) error {
	args := m.Called(ctx, userID, secret, recoveryCodes)
	return args.Error(0)
}

func (m *MockTwoFactorService) DisableTwoFactor(ctx context.Context, userID uuid.UUID) error {
	args := m.Called(ctx, userID)
	return args.Error(0)
}

// MockUserService mocks UserService
type MockUserService struct {
	mock.Mock
}

func (m *MockUserService) GetByID(userID uuid.UUID) (*models.User, error) {
	args := m.Called(userID)
	if args.Get(0) == nil {
		return nil, args.Error(1)
	}
	return args.Get(0).(*models.User), args.Error(1)
}

func setupTestTwoFactorRouter(mockTwoFactorService *MockTwoFactorService, mockUserService *MockUserService) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()

	logger := zap.NewNop()
	handler := NewTwoFactorHandlerWithInterface(mockTwoFactorService, mockUserService, logger)

	api := router.Group("/api/v1/auth/2fa")
	api.Use(func(c *gin.Context) {
		userIDStr := c.GetHeader("X-User-ID")
		if userIDStr != "" {
			uid, err := uuid.Parse(userIDStr)
			if err == nil {
				c.Set("user_id", uid)
			}
		}
		c.Next()
	})
	{
		api.POST("/setup", handler.SetupTwoFactor)
		api.POST("/verify", handler.VerifyTwoFactor)
		api.POST("/disable", handler.DisableTwoFactor)
		api.GET("/status", handler.GetTwoFactorStatus)
	}

	return router
}

func TestTwoFactorHandler_SetupTwoFactor_Success(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()
	mockUser := &models.User{
		ID:       userID,
		Email:    "test@example.com",
		Username: "testuser",
	}
	mockSetup := &services.TwoFactorSetup{
		Secret:        "TEST_SECRET",
		QRCodeURL:     "otpauth://totp/Veza:test@example.com?secret=TEST_SECRET",
		RecoveryCodes: []string{"CODE1", "CODE2"},
	}

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(false, nil)
	mockUserService.On("GetByID", userID).Return(mockUser, nil)
	mockTwoFactorService.On("GenerateSecret", mockUser).Return(mockSetup, nil)

	// Execute
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/setup", nil)
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusOK, w.Code)
	var response map[string]interface{}
	err := json.Unmarshal(w.Body.Bytes(), &response)
	assert.NoError(t, err)
	assert.True(t, response["success"].(bool))
	mockTwoFactorService.AssertExpectations(t)
	mockUserService.AssertExpectations(t)
}

func TestTwoFactorHandler_SetupTwoFactor_AlreadyEnabled(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(true, nil)

	// Execute
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/setup", nil)
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusBadRequest, w.Code)
	mockTwoFactorService.AssertNotCalled(t, "GenerateSecret")
	mockUserService.AssertNotCalled(t, "GetByID")
}

func TestTwoFactorHandler_SetupTwoFactor_Unauthorized(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	// Execute - No X-User-ID header
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/setup", nil)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.True(t, w.Code == http.StatusUnauthorized || w.Code == http.StatusForbidden)
	mockTwoFactorService.AssertNotCalled(t, "GetTwoFactorStatus")
}

func TestTwoFactorHandler_VerifyTwoFactor_Success(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()
	reqBody := VerifyTwoFactorRequest{
		Secret: "TEST_SECRET",
		Code:   "123456",
	}
	recoveryCodes := []string{"CODE1", "CODE2"}

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(false, nil)
	mockTwoFactorService.On("VerifyTOTPCode", "TEST_SECRET", "123456").Return(true)
	mockTwoFactorService.On("GenerateRecoveryCodes").Return(recoveryCodes)
	mockTwoFactorService.On("EnableTwoFactor", mock.Anything, userID, "TEST_SECRET", recoveryCodes).Return(nil)

	body, _ := json.Marshal(reqBody)

	// Execute
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/verify", bytes.NewBuffer(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusOK, w.Code)
	mockTwoFactorService.AssertExpectations(t)
}

func TestTwoFactorHandler_VerifyTwoFactor_InvalidCode(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()
	reqBody := VerifyTwoFactorRequest{
		Secret: "TEST_SECRET",
		Code:   "000000",
	}

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(false, nil)
	mockTwoFactorService.On("VerifyTOTPCode", "TEST_SECRET", "000000").Return(false)

	body, _ := json.Marshal(reqBody)

	// Execute
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/verify", bytes.NewBuffer(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusBadRequest, w.Code)
	mockTwoFactorService.AssertNotCalled(t, "EnableTwoFactor")
}

func TestTwoFactorHandler_DisableTwoFactor_Success(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()
	reqBody := DisableTwoFactorRequest{
		Password: "password123",
	}

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(true, nil)
	mockTwoFactorService.On("DisableTwoFactor", mock.Anything, userID).Return(nil)

	body, _ := json.Marshal(reqBody)

	// Execute
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/disable", bytes.NewBuffer(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusOK, w.Code)
	mockTwoFactorService.AssertExpectations(t)
}

func TestTwoFactorHandler_DisableTwoFactor_NotEnabled(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()
	reqBody := DisableTwoFactorRequest{
		Password: "password123",
	}

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(false, nil)

	body, _ := json.Marshal(reqBody)

	// Execute
	req, _ := http.NewRequest("POST", "/api/v1/auth/2fa/disable", bytes.NewBuffer(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusBadRequest, w.Code)
	mockTwoFactorService.AssertNotCalled(t, "DisableTwoFactor")
}

func TestTwoFactorHandler_GetTwoFactorStatus_Success(t *testing.T) {
	// Setup
	mockTwoFactorService := new(MockTwoFactorService)
	mockUserService := new(MockUserService)
	router := setupTestTwoFactorRouter(mockTwoFactorService, mockUserService)

	userID := uuid.New()

	mockTwoFactorService.On("GetTwoFactorStatus", mock.Anything, userID).Return(true, nil)

	// Execute
	req, _ := http.NewRequest("GET", "/api/v1/auth/2fa/status", nil)
	req.Header.Set("X-User-ID", userID.String())
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	// Assert
	assert.Equal(t, http.StatusOK, w.Code)
	var response map[string]interface{}
	err := json.Unmarshal(w.Body.Bytes(), &response)
	assert.NoError(t, err)
	assert.True(t, response["success"].(bool))
	data := response["data"].(map[string]interface{})
	assert.True(t, data["enabled"].(bool))
	mockTwoFactorService.AssertExpectations(t)
}