"""Tests for application/api/answer/routes/stream.py""" import json import uuid from unittest.mock import MagicMock, patch import pytest # Static IDs — valid 24-hex-char strings _CONV_ID = "507f1f77bcf86cd799439011" _AGENT_ID = "507f1f77bcf86cd799439012" @pytest.fixture def mock_stream_processor(): """Create a mock StreamProcessor for stream tests.""" with patch( "application.api.answer.routes.stream.StreamProcessor" ) as MockProcessor: processor = MagicMock() processor.decoded_token = {"sub": "test_user"} processor.conversation_id = _CONV_ID processor.agent_config = {} processor.agent_id = _AGENT_ID processor.is_shared_usage = False processor.shared_token = None processor.model_id = "gpt-4" processor.build_agent.return_value = MagicMock() MockProcessor.return_value = processor yield processor @pytest.fixture def stream_client(mock_mongo_db, flask_app): """Create a test client with the stream route registered.""" from flask_restx import Api from application.api.answer.routes.stream import answer_ns api = Api(flask_app) api.add_namespace(answer_ns) flask_app.config["TESTING"] = True return flask_app.test_client() @pytest.mark.unit class TestStreamResourcePost: def test_missing_question_returns_400(self, stream_client, mock_stream_processor): resp = stream_client.post( "/stream", data=json.dumps({}), content_type="application/json", ) assert resp.status_code == 400 def test_successful_stream(self, stream_client, mock_stream_processor): def fake_stream(*args, **kwargs): yield f'data: {json.dumps({"type": "answer", "answer": "Hi"})}\n\n' yield f'data: {json.dumps({"type": "end"})}\n\n' with patch( "application.api.answer.routes.stream.StreamResource.validate_request", return_value=None, ), patch( "application.api.answer.routes.stream.StreamResource.check_usage", return_value=None, ), patch( "application.api.answer.routes.stream.StreamResource.complete_stream", side_effect=fake_stream, ): resp = stream_client.post( "/stream", data=json.dumps({"question": "What is Python?"}), content_type="application/json", ) assert resp.status_code == 200 assert "text/event-stream" in resp.content_type data = resp.get_data(as_text=True) assert '"type": "answer"' in data assert '"answer": "Hi"' in data def test_unauthorized_returns_401_stream( self, stream_client, mock_stream_processor ): mock_stream_processor.decoded_token = None with patch( "application.api.answer.routes.stream.StreamResource.validate_request", return_value=None, ): resp = stream_client.post( "/stream", data=json.dumps({"question": "test"}), content_type="application/json", ) assert resp.status_code == 401 assert "text/event-stream" in resp.content_type data = resp.get_data(as_text=True) assert "Unauthorized" in data def test_usage_exceeded_returns_error( self, stream_client, mock_stream_processor ): with patch( "application.api.answer.routes.stream.StreamResource.validate_request", return_value=None, ), patch( "application.api.answer.routes.stream.StreamResource.check_usage", ) as mock_check: mock_check.return_value = ({"error": "Usage limit exceeded"}, 429) resp = stream_client.post( "/stream", data=json.dumps({"question": "test"}), content_type="application/json", ) assert resp.status_code == 429 def test_value_error_returns_400_stream( self, stream_client, mock_stream_processor ): mock_stream_processor.build_agent.side_effect = ValueError("bad data") with patch( "application.api.answer.routes.stream.StreamResource.validate_request", return_value=None, ): resp = stream_client.post( "/stream", data=json.dumps({"question": "test"}), content_type="application/json", ) assert resp.status_code == 400 assert "text/event-stream" in resp.content_type data = resp.get_data(as_text=True) assert "Malformed request body" in data def test_general_exception_returns_400_stream( self, stream_client, mock_stream_processor ): mock_stream_processor.build_agent.side_effect = RuntimeError("crash") with patch( "application.api.answer.routes.stream.StreamResource.validate_request", return_value=None, ): resp = stream_client.post( "/stream", data=json.dumps({"question": "test"}), content_type="application/json", ) assert resp.status_code == 400 assert "text/event-stream" in resp.content_type data = resp.get_data(as_text=True) assert "Unknown error occurred" in data def test_index_in_data_requires_conversation_id( self, stream_client, mock_stream_processor ): """When 'index' is present, validate_request is called with require_conversation_id=True.""" resp = stream_client.post( "/stream", data=json.dumps({"question": "test", "index": 0}), content_type="application/json", ) # Should get 400 since conversation_id is missing assert resp.status_code == 400 def test_stream_passes_attachments_and_index( self, stream_client, mock_stream_processor ): """Verify attachments and index params are forwarded to complete_stream.""" def fake_stream(*args, **kwargs): yield f'data: {json.dumps({"type": "end"})}\n\n' conv_id = str(uuid.uuid4()) with patch( "application.api.answer.routes.stream.StreamResource.validate_request", return_value=None, ), patch( "application.api.answer.routes.stream.StreamResource.check_usage", return_value=None, ), patch( "application.api.answer.routes.stream.StreamResource.complete_stream", side_effect=fake_stream, ) as mock_complete: resp = stream_client.post( "/stream", data=json.dumps( { "question": "test", "conversation_id": conv_id, "index": 3, "attachments": ["att1", "att2"], } ), content_type="application/json", ) assert resp.status_code == 200 call_kwargs = mock_complete.call_args assert call_kwargs.kwargs.get("index") == 3 assert call_kwargs.kwargs.get("attachment_ids") == ["att1", "att2"]