feat: Add unit tests for /create-newsletter endpoint

Author: BhaveshKukreja29

docker/init-db.sh

@@ -9,6 +9,7 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-E
     CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
 
     -- You can add other initialization SQL here
+    CREATE DATABASE mxtoaitest;
 
     GRANT ALL PRIVILEGES ON DATABASE $POSTGRES_DB TO $POSTGRES_USER;
 EOSQL

docker/worker.dockerfile

@@ -26,6 +26,9 @@ RUN poetry config virtualenvs.create false && poetry install --no-root --no-inte
 # Copy application code
 COPY mxgo ./mxgo
 
+# Copy test code into the container
+COPY ./tests /app/tests 
+
 # Create directories
 RUN mkdir -p /app/attachments
 

mxgo/tools/scheduled_tasks_tool.py

@@ -10,6 +10,8 @@
 from mxgo.config import SCHEDULED_TASKS_MINIMUM_INTERVAL_HOURS
 from mxgo.crud import create_task, delete_task, update_task_status
 from mxgo.db import init_db_connection
+from mxgo.email_handles import DEFAULT_EMAIL_HANDLES
+from mxgo.instruction_resolver import ProcessingInstructionsResolver
 from mxgo.models import TaskStatus
 from mxgo.request_context import RequestContext
 from mxgo.scheduling.scheduled_task_executor import execute_scheduled_task
@@ -247,11 +249,23 @@ def forward(
             }
 
         try:
+            # Resolve the handle alias before validation
+            resolved_handle_alias = future_handle_alias
+            if future_handle_alias:
+                try:
+                    # Resolve the provided alias (e.g., 'remind') to its canonical handle (e.g., 'schedule')
+                    resolver = ProcessingInstructionsResolver(DEFAULT_EMAIL_HANDLES)
+                    resolved_instructions = resolver(future_handle_alias)
+                    resolved_handle_alias = resolved_instructions.handle
+                    logger.info(f"Resolved future handle alias '{future_handle_alias}' to '{resolved_handle_alias}'")
+                except Exception:
+                    logger.warning(f"Could not resolve handle alias '{future_handle_alias}'. Validation might fail.")
+
             # Validate input using Pydantic
             input_data = ScheduledTaskInput(
                 cron_expression=cron_expression,
                 distilled_future_task_instructions=distilled_future_task_instructions,
-                future_handle_alias=future_handle_alias,
+                future_handle_alias=resolved_handle_alias,
                 start_time=start_time,
                 end_time=end_time,
             )

poetry.lock

@@ -15,6 +15,7 @@
 import mxgo.validators
 from mxgo._logging import get_logger
 from mxgo.api import app
+from mxgo.config import NEWSLETTER_LIMITS_BY_PLAN
 from mxgo.schemas import EmailSuggestionResponse, SuggestionDetail, UserPlan
 from tests.generate_test_jwt import generate_test_jwt
 
@@ -1207,3 +1208,119 @@ def test_user_info_api_different_user_emails(
     # Verify the correct email was used for lookups
     mock_get_user_plan.assert_called_once_with("[email protected]")
     mock_get_customer_id.assert_called_once_with("[email protected]")
+
+
+@patch("mxgo.auth.JWT_SECRET", "test_secret_key_for_development_only")
+class TestCreateNewsletter:
+    @pytest.fixture
+    def mock_dependencies(self):
+        """Mocks all external dependencies for the newsletter endpoint."""
+        with patch("mxgo.api.user.get_user_plan", new_callable=AsyncMock) as mock_get_plan, \
+             patch("mxgo.api.crud.count_active_tasks_for_user") as mock_count_tasks, \
+             patch("mxgo.api.whitelist.is_email_whitelisted", new_callable=AsyncMock) as mock_is_whitelisted, \
+             patch("mxgo.api.Scheduler.add_job") as mock_add_job, \
+             patch("mxgo.api.process_email_task.send") as mock_send_task:
+            
+            # Default happy path mocks
+            mock_get_plan.return_value = UserPlan.BETA
+            mock_count_tasks.return_value = 0
+            mock_is_whitelisted.return_value = (True, True) # Assume user is whitelisted
+            
+            yield {
+                "get_plan": mock_get_plan,
+                "count_tasks": mock_count_tasks,
+                "is_whitelisted": mock_is_whitelisted,
+                "add_job": mock_add_job,
+                "send_task": mock_send_task,
+            }
+
+    def test_create_newsletter_success_whitelisted(self, mock_dependencies):
+        """Test successful newsletter creation for a whitelisted BETA user."""
+        jwt_token = generate_test_jwt(email="[email protected]", user_id="test_user_123")
+
+        response = client.post(
+            "/create-newsletter",
+            headers={"Authorization": f"Bearer {jwt_token}"},
+            json={
+                "prompt": "Weekly AI news",
+                "schedule": {
+                    "type": "RECURRING_WEEKLY",
+                    "recurring_weekly": {"days": ["friday"], "time": "10:00"},
+                },
+            },
+        )
+        assert response.status_code == 200
+        data = response.json()
+        assert data["is_scheduled"] is True
+        assert data["is_whitelisted"] is True
+        assert data["sample_email_sent"] is True
+        assert len(data["scheduled_task_ids"]) == 1
+        mock_dependencies["send_task"].assert_called_once()
+
+    def test_create_newsletter_task_limit_exceeded(self, mock_dependencies):
+        """Test that newsletter creation fails if the task limit is reached."""
+        mock_dependencies["get_plan"].return_value = UserPlan.BETA
+
+        jwt_token = generate_test_jwt(email="[email protected]", user_id="test_user_123")
+
+        # Set current tasks to the max allowed for BETA plan using the config variable
+        mock_dependencies["count_tasks"].return_value = NEWSLETTER_LIMITS_BY_PLAN[UserPlan.BETA]["max_tasks"]
+
+        response = client.post(
+            "/create-newsletter",
+            headers={"Authorization": f"Bearer {jwt_token}"},
+            json={
+                "prompt": "Another newsletter",
+                "schedule": {
+                    "type": "RECURRING_WEEKLY",
+                    "recurring_weekly": {"days": ["monday"], "time": "08:00"},
+                },
+            },
+        )
+        assert response.status_code == 403
+        assert "Newsletter limit" in response.json()["detail"]
+
+    def test_create_newsletter_interval_too_frequent(self, mock_dependencies):
+        """Test that creation fails if the cron interval is too short for the user's plan."""
+        mock_dependencies["get_plan"].return_value = UserPlan.BETA
+
+        jwt_token = generate_test_jwt(email="[email protected]", user_id="test_user_123")
+
+        response = client.post(
+            "/create-newsletter",
+            headers={"Authorization": f"Bearer {jwt_token}"},
+            json={
+                "prompt": "Daily newsletter",
+                "schedule": { 
+                    "type": "RECURRING_WEEKLY",
+                    "recurring_weekly": {"days": ["monday", "tuesday"], "time": "07:00"},
+                },
+            },
+        )
+        assert response.status_code == 400
+        assert "Cron interval is too frequent" in response.json()["detail"]
+        
+    def test_create_newsletter_not_whitelisted(self, mock_dependencies):
+        """Test behavior for a non-whitelisted user."""
+        mock_dependencies["is_whitelisted"].return_value = (False, False)
+        jwt_token = generate_test_jwt(email="[email protected]", user_id="test_user_123")
+
+        with patch("mxgo.api.whitelist.trigger_automatic_verification", new_callable=AsyncMock) as mock_trigger_verify:
+            response = client.post(
+                "/create-newsletter",
+                headers={"Authorization": f"Bearer {jwt_token}"},
+                json={
+                    "prompt": "A test newsletter",
+                    "schedule": {
+                        "type": "RECURRING_WEEKLY",
+                        "recurring_weekly": {"days": ["saturday"], "time": "12:00"},
+                    },
+                },
+            )
+            assert response.status_code == 200
+            data = response.json()
+            assert data["is_scheduled"] is True
+            assert data["is_whitelisted"] is False
+            assert data["sample_email_sent"] is False
+            mock_dependencies["send_task"].assert_not_called()
+            mock_trigger_verify.assert_called_once()