Fix some tests and add more

Author: satwikkansal

tests/test_api.py

@@ -9,7 +9,6 @@
 
 # Corrected import for FakeAsyncRedis based on common usage and docs
 from fakeredis import FakeAsyncRedis
-from fastapi import Response
 from fastapi.testclient import TestClient
 from freezegun import freeze_time  # For controlling time in tests
 
@@ -540,17 +539,16 @@ def assert_suggestions_error_response(response, expected_status=422):
 
 @patch.dict(os.environ, {"SUGGESTIONS_API_KEY": "valid-suggestions-key"})
 @patch("mxgo.auth.JWT_SECRET", "test_secret_key_for_development_only")
-@patch("mxgo.api.validate_email_whitelist", new_callable=AsyncMock)
 @patch("mxgo.api.generate_suggestions", new_callable=AsyncMock)
-def test_suggestions_api_success_single_request(mock_generate_suggestions, mock_validate_whitelist):
+def test_suggestions_api_success_single_request(mock_generate_suggestions):
     """Test successful suggestions API call with single request."""
-    mock_validate_whitelist.return_value = None  # User is whitelisted
 
     # Mock the suggestions response
 
     mock_response = EmailSuggestionResponse(
         email_identified="test-email-123",
         user_email_id="test@example.com",
+        overview="Test email requesting content summarization",
         suggestions=[
             SuggestionDetail(
                 suggestion_title="Summarize content",
@@ -574,7 +572,6 @@ def test_suggestions_api_success_single_request(mock_generate_suggestions, mock_
     response = make_suggestions_post_request(request_data)
 
     assert_suggestions_successful_response(response, expected_num_requests=1)
-    mock_validate_whitelist.assert_called_once()
     mock_generate_suggestions.assert_called_once()
 
     # Verify the response content
@@ -587,17 +584,16 @@ def test_suggestions_api_success_single_request(mock_generate_suggestions, mock_
 
 @patch.dict(os.environ, {"SUGGESTIONS_API_KEY": "valid-suggestions-key"})
 @patch("mxgo.auth.JWT_SECRET", "test_secret_key_for_development_only")
-@patch("mxgo.api.validate_email_whitelist", new_callable=AsyncMock)
 @patch("mxgo.api.generate_suggestions", new_callable=AsyncMock)
-def test_suggestions_api_success_multiple_requests(mock_generate_suggestions, mock_validate_whitelist):
+def test_suggestions_api_success_multiple_requests(mock_generate_suggestions):
     """Test successful suggestions API call with multiple requests."""
-    mock_validate_whitelist.return_value = None
 
     # Mock responses for each request
     mock_responses = [
         EmailSuggestionResponse(
             email_identified=f"test-email-{i}",
             user_email_id="test@example.com",
+            overview=f"Test email {i} overview",
             suggestions=[
                 SuggestionDetail(
                     suggestion_title="Ask anything",
@@ -629,7 +625,6 @@ def test_suggestions_api_success_multiple_requests(mock_generate_suggestions, mo
     response = make_suggestions_post_request(request_data)
 
     assert_suggestions_successful_response(response, expected_num_requests=3)
-    assert mock_validate_whitelist.call_count == 3
     assert mock_generate_suggestions.call_count == 3
 
 
@@ -655,22 +650,34 @@ def test_suggestions_api_invalid_jwt_token():
 
 @patch.dict(os.environ, {"SUGGESTIONS_API_KEY": "valid-suggestions-key"})
 @patch("mxgo.auth.JWT_SECRET", "test_secret_key_for_development_only")
-@patch("mxgo.api.validate_email_whitelist", new_callable=AsyncMock)
-def test_suggestions_api_user_not_whitelisted(mock_validate_whitelist):
-    """Test suggestions API when user is not whitelisted."""
-    # Mock whitelist validation to return an error response
-    mock_validate_whitelist.return_value = Response(
-        status_code=403, content=json.dumps({"detail": "User not whitelisted"})
+@patch("mxgo.api.generate_suggestions", new_callable=AsyncMock)
+def test_suggestions_api_user_not_whitelisted(mock_generate_suggestions):
+    """Test suggestions API processes normally since whitelist check was removed."""
+    # Mock successful suggestions generation
+    mock_response = EmailSuggestionResponse(
+        email_identified="test-email-123",
+        user_email_id="test@example.com",
+        overview="Test email processed without whitelist check",
+        suggestions=[
+            SuggestionDetail(
+                suggestion_title="Ask anything",
+                suggestion_id="suggest-1",
+                suggestion_to_email="ask@mxgo.ai",
+                suggestion_cc_emails=[],
+                suggestion_email_instructions="",
+            ),
+        ],
     )
+    mock_generate_suggestions.return_value = mock_response
 
     request_data = prepare_suggestions_request_data()
     response = make_suggestions_post_request(request_data)
 
-    # Should return 403 Forbidden instead of 200 with error suggestions
-    assert response.status_code == 403, "Should return 403 Forbidden for non-whitelisted user"
+    # Should return 200 OK since whitelist validation was removed
+    assert response.status_code == 200, "Should return 200 OK since whitelist check removed"
     response_json = response.json()
-    assert "detail" in response_json
-    assert "Email verification required" in response_json["detail"]
+    assert len(response_json) == 1
+    assert response_json[0]["email_identified"] == "test-email-123"
 
 
 @patch.dict(os.environ, {"SUGGESTIONS_API_KEY": "valid-suggestions-key"})
@@ -728,6 +735,7 @@ def test_suggestions_api_with_attachments(mock_generate_suggestions, mock_valida
     mock_response = EmailSuggestionResponse(
         email_identified="test-email-123",
         user_email_id="test@example.com",
+        overview="Email with PDF and Excel attachments requiring document analysis",
         suggestions=[
             SuggestionDetail(
                 suggestion_title="Summarize documents",
@@ -778,6 +786,7 @@ def test_suggestions_api_with_cc_emails(mock_generate_suggestions, mock_validate
     mock_response = EmailSuggestionResponse(
         email_identified="test-email-123",
         user_email_id="test@example.com",
+        overview="Meeting request with CC recipients",
         suggestions=[
             SuggestionDetail(
                 suggestion_title="Schedule meeting",
@@ -824,6 +833,7 @@ def test_suggestions_api_rate_limiting_integration(client_with_patched_redis):
         mock_response = EmailSuggestionResponse(
             email_identified="test-email-123",
             user_email_id="test@example.com",
+            overview="Test email for rate limiting integration",
             suggestions=[
                 SuggestionDetail(
                     suggestion_title="Ask anything",
@@ -851,6 +861,7 @@ def test_suggestions_api_subject_field_alias(mock_generate_suggestions, mock_val
     mock_response = EmailSuggestionResponse(
         email_identified="test-email-123",
         user_email_id="test@example.com",
+        overview="Test email with Subject field alias",
         suggestions=[
             SuggestionDetail(
                 suggestion_title="Ask anything",
@@ -886,6 +897,7 @@ def test_suggestions_api_default_suggestions_always_included(mock_generate_sugge
     mock_response = EmailSuggestionResponse(
         email_identified="test-email-123",
         user_email_id="test@example.com",
+        overview="Email with statistical claims requiring fact verification",
         suggestions=[
             SuggestionDetail(
                 suggestion_title="Fact check claims",

tests/test_api_integration.py

@@ -260,11 +260,25 @@ def mock_suggestions_model(self):
     @pytest.fixture
     def mock_generate_suggestions(self):
         """Mock generate_suggestions function."""
+        from mxgo.schemas import EmailSuggestionResponse, SuggestionDetail
+
         with patch("mxgo.api.generate_suggestions") as mock_gen:
-            mock_gen.return_value = AsyncMock()
-            mock_gen.return_value.email_identified = "test_email_123"
-            mock_gen.return_value.user_email_id = "test@example.com"
-            mock_gen.return_value.suggestions = []
+            mock_response = EmailSuggestionResponse(
+                email_identified="test_email_123",
+                user_email_id="test@example.com",
+                overview="Test email processed successfully",
+                suggestions=[
+                    SuggestionDetail(
+                        suggestion_title="Ask anything",
+                        suggestion_id="suggest-1",
+                        suggestion_to_email="ask@mxgo.ai",
+                        suggestion_cc_emails=[],
+                        suggestion_email_instructions="",
+                    )
+                ],
+                risk_analysis=None,
+            )
+            mock_gen.return_value = mock_response
             yield mock_gen
 
     @pytest.fixture
@@ -341,18 +355,20 @@ def test_suggestions_invalid_jwt_token(self, client, mock_env_vars, sample_sugge
         assert response.status_code == 401
 
     def test_suggestions_user_not_whitelisted(
-        self, client, mock_env_vars, valid_jwt_token, sample_suggestion_request, mock_suggestions_model
+        self, client, mock_env_vars, valid_jwt_token, sample_suggestion_request, mock_suggestions_model, mock_generate_suggestions
     ):
-        """Test /suggestions endpoint handles non-whitelisted users."""
+        """Test /suggestions endpoint processes normally since whitelist check was removed."""
         with patch("mxgo.validators.is_email_whitelisted", return_value=(False, False)):
             response = client.post(
                 "/suggestions",
                 json=sample_suggestion_request,
                 headers={"Authorization": f"Bearer {valid_jwt_token}"},
             )
 
-            assert response.status_code == 403
-            assert "Email verification required" in response.json()["detail"]
+            assert response.status_code == 200
+            response_json = response.json()
+            assert len(response_json) == 1
+            assert response_json[0]["email_identified"] == "test_email_123"
 
 
 class TestBackwardCompatibility:

tests/test_process_email.py

@@ -24,6 +24,9 @@
 from mxgo.tasks import process_email_task
 from mxgo.tools.pdf_export_tool import MAX_FILENAME_LENGTH, PDFExportTool
 
+# Import email_agent to ensure TokenUsage monkey patch is applied
+import mxgo.agents.email_agent  # noqa: F401
+
 AttachmentFileContent = tuple[str, bytes, str]  # (filename, content_bytes, content_type)
 
 

tests/test_suggestions.py

@@ -429,3 +429,165 @@ async def test_suggestions_integration_error_handling(prepare_suggestion_request
     assert len(response.suggestions) >= 1
     has_ask = any(s.suggestion_to_email == "ask@mxgo.ai" for s in response.suggestions)
     assert has_ask, "Should return default suggestion even for empty content"
+
+
+@pytest.mark.asyncio
+async def test_risk_analysis_functionality(prepare_suggestion_request_data):
+    """Test risk analysis functionality in suggestions."""
+    # Test suspicious email content
+    request_data = prepare_suggestion_request_data(
+        email_identified="risky-email-001",
+        user_email_id="victim@company.com",
+        sender_email="attacker@suspicious.com",
+        subject="URGENT: Verify Your Account Now!",
+        email_content="""
+        Dear Customer,
+
+        Your account has been compromised. Click here immediately to verify:
+        https://fake-bank.com/verify-account
+
+        You have 24 hours or your account will be closed permanently.
+
+        Best regards,
+        Security Team
+        """,
+        attachments=[
+            {"filename": "invoice.pdf.exe", "file_type": "application/x-executable", "file_size": 2048000}
+        ]
+    )
+
+    request_obj = EmailSuggestionRequest(
+        email_identified=request_data["email_identified"],
+        user_email_id=request_data["user_email_id"],
+        sender_email=request_data["sender_email"],
+        cc_emails=request_data["cc_emails"],
+        Subject=request_data["Subject"],
+        email_content=request_data["email_content"],
+        attachments=[EmailSuggestionAttachmentSummary(**att) for att in request_data["attachments"]],
+    )
+
+    model = get_suggestions_model()
+    response = await generate_suggestions(request_obj, model)
+
+    # Validate basic response structure
+    assert response.email_identified == request_data["email_identified"]
+    assert response.user_email_id == request_data["user_email_id"]
+    assert response.overview is not None
+    assert len(response.overview) > 0
+    assert len(response.suggestions) >= 1
+
+    # Validate risk analysis is present
+    assert response.risk_analysis is not None
+    assert hasattr(response.risk_analysis, "risk_prob_pct")
+    assert hasattr(response.risk_analysis, "spam_prob_pct")
+    assert hasattr(response.risk_analysis, "ai_likelihood_pct")
+    assert 0 <= response.risk_analysis.risk_prob_pct <= 100
+    assert 0 <= response.risk_analysis.spam_prob_pct <= 100
+    assert 0 <= response.risk_analysis.ai_likelihood_pct <= 100
+
+    # For suspicious email, expect higher risk scores
+    # (Note: actual scores depend on LLM, so we just check structure)
+    assert isinstance(response.risk_analysis.risk_reason, str)
+    assert isinstance(response.risk_analysis.spam_reason, str)
+    assert isinstance(response.risk_analysis.ai_explanation, str)
+
+
+@pytest.mark.asyncio
+async def test_risk_analysis_benign_email(prepare_suggestion_request_data):
+    """Test risk analysis with benign email content."""
+    request_data = prepare_suggestion_request_data(
+        email_identified="safe-email-001",
+        user_email_id="employee@company.com",
+        sender_email="colleague@company.com",
+        subject="Weekly team meeting notes",
+        email_content="""
+        Hi team,
+
+        Here are the notes from our weekly meeting:
+
+        1. Project Alpha is on track for Q4 delivery
+        2. New team member Sarah will start Monday
+        3. Budget review scheduled for next Friday
+
+        Let me know if you have any questions.
+
+        Best,
+        John
+        """,
+    )
+
+    request_obj = EmailSuggestionRequest(
+        email_identified=request_data["email_identified"],
+        user_email_id=request_data["user_email_id"],
+        sender_email=request_data["sender_email"],
+        cc_emails=request_data["cc_emails"],
+        Subject=request_data["Subject"],
+        email_content=request_data["email_content"],
+        attachments=[],
+    )
+
+    model = get_suggestions_model()
+    response = await generate_suggestions(request_obj, model)
+
+    # Validate risk analysis structure
+    assert response.risk_analysis is not None
+    assert 0 <= response.risk_analysis.risk_prob_pct <= 100
+    assert 0 <= response.risk_analysis.spam_prob_pct <= 100
+    assert 0 <= response.risk_analysis.ai_likelihood_pct <= 100
+
+    # All reason fields should be strings (may be empty)
+    assert isinstance(response.risk_analysis.risk_reason, str)
+    assert isinstance(response.risk_analysis.spam_reason, str)
+    assert isinstance(response.risk_analysis.ai_explanation, str)
+
+
+@pytest.mark.asyncio
+async def test_overview_generation(prepare_suggestion_request_data):
+    """Test that overview field is properly generated."""
+    request_data = prepare_suggestion_request_data(
+        email_identified="overview-test-001",
+        user_email_id="user@company.com",
+        sender_email="client@external.com",
+        subject="Quarterly Business Review Meeting Request",
+        email_content="""
+        Dear Team,
+
+        I would like to schedule our quarterly business review for next week.
+        Please let me know your availability for Tuesday or Wednesday afternoon.
+
+        We'll need to cover:
+        - Q3 performance analysis
+        - Budget planning for Q4
+        - New client onboarding process
+
+        Thanks,
+        Client
+        """,
+    )
+
+    request_obj = EmailSuggestionRequest(
+        email_identified=request_data["email_identified"],
+        user_email_id=request_data["user_email_id"],
+        sender_email=request_data["sender_email"],
+        cc_emails=request_data["cc_emails"],
+        Subject=request_data["Subject"],
+        email_content=request_data["email_content"],
+        attachments=[],
+    )
+
+    model = get_suggestions_model()
+    response = await generate_suggestions(request_obj, model)
+
+    # Overview should be non-empty and meaningful
+    assert response.overview is not None
+    assert len(response.overview.strip()) > 0
+    assert len(response.overview) <= 500  # Should be reasonably concise
+
+    # Overview should ideally relate to the subject/content
+    # (This is a basic check - actual content quality depends on LLM)
+    overview_lower = response.overview.lower()
+    request_data["Subject"].lower().split()
+
+    # At least some connection to the email content should exist
+    # (This is a weak test since LLM behavior can vary)
+    assert len(overview_lower) > 10  # Not just a few words