fix: Fix Swagger Schema gen within the two_factor_enforce logic (#38485)

posthog/conftest.py

@@ -201,9 +201,13 @@ def cache():
 
 
 @pytest.fixture(autouse=True)
-def mock_two_factor_sso_enforcement_check(mocker):
+def mock_two_factor_sso_enforcement_check(request, mocker):
     """
     Mock the two_factor_session.is_domain_sso_enforced check to return False for all tests.
+    Can be disabled by using @pytest.mark.no_mock_two_factor_sso_enforcement_check decorator.
     """
+    if "no_mock_two_factor_sso_enforcement_check" in request.keywords:
+        return
+
     mocker.patch("posthog.helpers.two_factor_session.is_domain_sso_enforced", return_value=False)
     mocker.patch("posthog.helpers.two_factor_session.is_sso_authentication_backend", return_value=False)

posthog/helpers/two_factor_session.py

@@ -146,6 +146,9 @@ def is_sso_authentication_backend(request: HttpRequest):
     SSO_AUTHENTICATION_BACKENDS = []
     NON_SSO_AUTHENTICATION_BACKENDS = ["axes.backends.AxesBackend", "django.contrib.auth.backends.ModelBackend"]
 
+    if not hasattr(request, "session"):
+        return False
+
     # Check if we're in EE, if yes, use the EE settings, otherwise use the posthog settings
     try:
         from ee import settings

posthog/test/test_two_factor_enforcement.py

@@ -1,6 +1,7 @@
 import time
 import datetime
 
+import pytest
 from unittest.mock import Mock, patch
 
 from django.conf import settings
@@ -512,3 +513,9 @@ class TestUserTwoFactorSessionIntegration(TestCase):
         mock_totp_form.assert_called_once_with("1234567890abcdef1234", self.user, data={"token": "123456"})
         mock_form_instance.save.assert_called_once()
         mock_send_email.delay.assert_called_once_with(self.user.id)
+
+    @pytest.mark.no_mock_two_factor_sso_enforcement_check
+    def test_doesnt_break_swagger_schema(self):
+        """Test that schema generation works without session middleware errors"""
+        response = self.client.get("/api/schema/")
+        self.assertEqual(response.status_code, 200)