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