diff --git a/llm_be/chat_backend/oauth.py b/llm_be/chat_backend/oauth.py index 92a3c77..4d3075a 100644 --- a/llm_be/chat_backend/oauth.py +++ b/llm_be/chat_backend/oauth.py @@ -338,6 +338,11 @@ def upsert_drive_connection( *, user: CustomUser, kind: str, profile: ProviderProfile ) -> DriveConnection: """Create/refresh a DriveConnection from a link_drive/link_company_drive callback (#47).""" + if not user.company_id: + raise OAuthError( + "no_company", + "A company is required before connecting a Drive account.", + ) expires_at = _token_expiry(profile.expires_in) lookup_user = user if kind == DriveConnection.Kind.PERSONAL else None connection = DriveConnection.objects.filter( diff --git a/llm_be/chat_backend/tests/test_oauth.py b/llm_be/chat_backend/tests/test_oauth.py index 3962a6d..d01d162 100644 --- a/llm_be/chat_backend/tests/test_oauth.py +++ b/llm_be/chat_backend/tests/test_oauth.py @@ -431,6 +431,25 @@ class OAuthCallbackDriveLinkTestCase(APITestCase): self.assertFalse(DriveConnection.objects.exists()) mock_exchange.assert_not_called() + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_rejects_user_without_company(self, mock_exchange): + """#55 — users with no company must not hit DriveConnection.company NOT NULL.""" + self.user.company = None + self.user.save(update_fields=["company"]) + mock_exchange.return_value = _google_profile() + state = self._state(intent="link_drive") + + response = self.client.get( + reverse("oauth_callback", kwargs={"provider": "google"}), + {"code": "auth-code", "state": state}, + ) + + self.assertEqual(response.status_code, status.HTTP_302_FOUND) + params = parse_qs(urlparse(response["Location"]).query) + self.assertEqual(params["error"], ["no_company"]) + self.assertFalse(DriveConnection.objects.exists()) + mock_exchange.assert_not_called() + def test_callback_missing_user_id_in_state_is_rejected(self): # Simulates a forged/legacy state payload without the linking user. from django.core import signing diff --git a/llm_be/chat_backend/views_oauth.py b/llm_be/chat_backend/views_oauth.py index de4104e..8434556 100644 --- a/llm_be/chat_backend/views_oauth.py +++ b/llm_be/chat_backend/views_oauth.py @@ -229,6 +229,11 @@ class OAuthCallbackView(APIView): "forbidden", "Only a company manager can connect a company Drive.", ) + if not user.company_id: + raise OAuthError( + "no_company", + "A company is required before connecting a Drive account.", + ) assert_feature_allowed(user, "rag") redirect_uri = _callback_redirect_uri(request, provider)