eigent/backend/tests/app/controller/test_chat_controller.py

# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========

import os
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import Response
from fastapi.responses import StreamingResponse
from fastapi.testclient import TestClient
from pydantic import ValidationError

from app.controller.chat_controller import (
    human_reply,
    improve,
    install_mcp,
    post,
    stop,
    supplement,
)
from app.exception.exception import UserException
from app.model.chat import Chat, HumanReply, McpServers, Status, SupplementChat


@pytest.mark.unit
class TestChatController:
    """Test cases for chat controller endpoints."""

    @pytest.mark.asyncio
    async def test_post_chat_endpoint_success(
        self,
        sample_chat_data,
        mock_request,
        mock_task_lock,
        mock_environment_variables,
    ):
        """Test successful chat initialization."""
        chat_data = Chat(**sample_chat_data)

        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller.step_solve"
            ) as mock_step_solve,
            patch("app.controller.chat_controller.load_dotenv"),
            patch("app.controller.chat_controller.set_current_task_id"),
            patch(
                "app.controller.chat_controller._prepare_browser_for_request_with_timeout",
                new=AsyncMock(return_value=True),
            ),
            patch("pathlib.Path.mkdir"),
            patch("pathlib.Path.home", return_value=MagicMock()),
        ):
            # Mock async generator
            async def mock_generator():
                yield "data: test_response\n\n"
                yield "data: test_response_2\n\n"

            mock_step_solve.return_value = mock_generator()

            response = await post(chat_data, mock_request)

            assert isinstance(response, StreamingResponse)
            assert response.media_type == "text/event-stream"

    @pytest.mark.asyncio
    async def test_post_chat_sets_run_context_and_third_party_env(
        self, sample_chat_data, mock_request, mock_task_lock
    ):
        """Run-scoped values stay in RunContext; CAMEL path keys are published."""
        chat_data = Chat(**sample_chat_data)

        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller.step_solve"
            ) as mock_step_solve,
            patch("app.controller.chat_controller.load_dotenv"),
            patch("app.controller.chat_controller.set_current_task_id"),
            patch(
                "app.controller.chat_controller._prepare_browser_for_request",
                return_value=True,
            ),
            patch("pathlib.Path.mkdir"),
            patch("pathlib.Path.home", return_value=MagicMock()),
            patch.dict(os.environ, {}, clear=True),
        ):

            async def mock_generator():
                yield "data: test_response\n\n"

            mock_step_solve.return_value = mock_generator()

            await post(chat_data, mock_request)

            run_context = mock_task_lock.run_context
            assert run_context.api_key == "test_key"
            assert run_context.api_base_url == "https://api.openai.com/v1"
            assert run_context.browser_port == 8080
            assert os.environ.get("CAMEL_LOG_DIR") == str(
                run_context.camel_log_dir
            )
            assert os.environ.get("CAMEL_WORKDIR") == str(
                run_context.task_output_root
            )

    @pytest.mark.asyncio
    async def test_post_chat_sets_cdp_url_when_browser_ready(
        self, sample_chat_data, mock_request, mock_task_lock
    ):
        """Web mode should set EIGENT_CDP_URL after successful browser ensure."""
        chat_data = Chat(**sample_chat_data)
        mock_request.state = SimpleNamespace()

        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller.step_solve"
            ) as mock_step_solve,
            patch(
                "app.controller.chat_controller.is_cdp_url_available",
                return_value=False,
            ),
            patch(
                "app.controller.chat_controller.ensure_cdp_browser_endpoint",
                return_value="http://127.0.0.1:8080",
            ),
            patch("app.controller.chat_controller.load_dotenv"),
            patch("app.controller.chat_controller.set_current_task_id"),
            patch("pathlib.Path.mkdir"),
            patch("pathlib.Path.home", return_value=MagicMock()),
            patch.dict(os.environ, {}, clear=True),
        ):

            async def mock_generator():
                yield "data: test_response\n\n"

            mock_step_solve.return_value = mock_generator()

            await post(chat_data, mock_request)

            assert (
                mock_task_lock.run_context.cdp_url == "http://127.0.0.1:8080"
            )
            assert mock_task_lock.run_context.browser_port == 8080
            assert mock_request.state.browser_available is True

    @pytest.mark.asyncio
    async def test_post_chat_clears_cdp_url_when_browser_unavailable(
        self, sample_chat_data, mock_request, mock_task_lock
    ):
        """Web mode should mark browser unavailable and clear EIGENT_CDP_URL."""
        chat_data = Chat(**sample_chat_data)
        mock_request.state = SimpleNamespace()

        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller.step_solve"
            ) as mock_step_solve,
            patch(
                "app.controller.chat_controller.is_cdp_url_available",
                return_value=False,
            ),
            patch(
                "app.controller.chat_controller.ensure_cdp_browser_endpoint",
                return_value=None,
            ),
            patch("app.controller.chat_controller.load_dotenv"),
            patch("app.controller.chat_controller.set_current_task_id"),
            patch("pathlib.Path.mkdir"),
            patch("pathlib.Path.home", return_value=MagicMock()),
            patch.dict(
                os.environ,
                {"EIGENT_CDP_URL": "http://127.0.0.1:9222"},
                clear=True,
            ),
        ):

            async def mock_generator():
                yield "data: test_response\n\n"

            mock_step_solve.return_value = mock_generator()

            await post(chat_data, mock_request)

            assert mock_task_lock.run_context.cdp_url is None
            assert mock_task_lock.run_context.browser_port == 8080
            assert mock_request.state.browser_available is False

    @pytest.mark.asyncio
    async def test_post_chat_preserves_existing_cdp_url(
        self, sample_chat_data, mock_request, mock_task_lock
    ):
        chat_data = Chat(**sample_chat_data)
        mock_request.state = SimpleNamespace()

        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller.step_solve"
            ) as mock_step_solve,
            patch(
                "app.controller.chat_controller.is_cdp_url_available",
                return_value=True,
            ),
            patch(
                "app.controller.chat_controller.ensure_cdp_browser_endpoint",
            ) as mock_ensure_browser,
            patch("app.controller.chat_controller.load_dotenv"),
            patch("app.controller.chat_controller.set_current_task_id"),
            patch("pathlib.Path.mkdir"),
            patch("pathlib.Path.home", return_value=MagicMock()),
            patch.dict(
                os.environ,
                {"EIGENT_CDP_URL": "http://worker-17:9222"},
                clear=True,
            ),
        ):

            async def mock_generator():
                yield "data: test_response\n\n"

            mock_step_solve.return_value = mock_generator()

            await post(chat_data, mock_request)

            assert (
                mock_task_lock.run_context.cdp_url == "http://worker-17:9222"
            )
            assert mock_task_lock.run_context.browser_port == 9222
            assert mock_request.state.browser_available is True
            mock_ensure_browser.assert_not_called()

    @pytest.mark.asyncio
    async def test_improve_chat_success(self, mock_task_lock, mock_request):
        """Test successful chat improvement."""
        task_id = "test_task_123"
        supplement_data = SupplementChat(question="Improve this code")
        mock_task_lock.status = Status.processing

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller._prepare_browser_for_request_with_timeout",
                new=AsyncMock(return_value=True),
            ),
        ):
            response = await improve(task_id, supplement_data, mock_request)

            assert isinstance(response, Response)
            assert response.status_code == 201
            mock_task_lock.put_queue.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_improve_chat_task_done_resets_to_confirming(
        self, mock_task_lock, mock_request
    ):
        """Test improvement when task is done resets status to confirming."""
        task_id = "test_task_123"
        supplement_data = SupplementChat(question="Improve this code")
        mock_task_lock.status = Status.done

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller._prepare_browser_for_request_with_timeout",
                new=AsyncMock(return_value=True),
            ),
        ):
            response = await improve(task_id, supplement_data, mock_request)

            assert mock_task_lock.status == Status.confirming
            assert isinstance(response, Response)
            assert response.status_code == 201

    @pytest.mark.asyncio
    async def test_improve_skips_durable_on_run_start_when_rotation_failed(
        self, mock_task_lock, mock_request
    ):
        """R27-1 regression: if run_context did not rotate to data.task_id
        (e.g. freeze_task_directories_for raised inside the swallowing
        try/except), we must NOT call MemoryService.on_run_start against the
        previous, already-finalized run id -- doing so would reset its
        status.json back to "running" and the finalize dedup set then blocks
        future end-of-turn writers from closing it again.
        """

        # Use a real RunContext-shaped stand-in so getattr access works.
        stale_run = SimpleNamespace(
            run_id="stale_run_from_previous_turn",
            space_id="blank_space_x",
            project_id="project_x",
            task_id="stale_task_id",
            browser_port=9222,
            cdp_url=None,
            user_id="42",
            email="u@example.com",
        )
        # Make isinstance(stale_run, RunContext) return True via patching.
        mock_task_lock.run_context = stale_run
        mock_task_lock.status = Status.processing
        mock_task_lock.email = "u@example.com"

        memory_service = MagicMock()
        supplement_data = SupplementChat(
            question="next turn", task_id="brand_new_task_id"
        )

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller.RunContext",
                new=SimpleNamespace,
            ),
            patch(
                "app.controller.chat_controller.get_workspace_resolver",
                side_effect=RuntimeError(
                    "simulated resolver failure -- rotation never happens"
                ),
            ),
            patch(
                "app.controller.chat_controller.get_memory_service",
                return_value=memory_service,
            ),
            patch(
                "app.controller.chat_controller._prepare_browser_for_request_with_timeout",
                new=AsyncMock(return_value=True),
            ),
        ):
            response = await improve(
                "project_x", supplement_data, mock_request
            )

            assert isinstance(response, Response)
            # The improve request itself still succeeds -- chat must not break
            # because durable memory is unhappy.
            assert response.status_code == 201
            # Critical assertion: on_run_start was NOT called against the stale
            # context. The R26 fix only checked data.task_id; R27 strengthens
            # it to compare refreshed_context.run_id == data.task_id.
            memory_service.on_run_start.assert_not_called()

    def test_supplement_chat_success(self, mock_task_lock):
        """Test successful chat supplementation."""
        task_id = "test_task_123"
        supplement_data = SupplementChat(question="Add more details")
        mock_task_lock.status = Status.done

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller._queue_action_from_worker"
            ) as mock_queue_action,
        ):
            response = supplement(task_id, supplement_data)

            assert isinstance(response, Response)
            assert response.status_code == 201
            mock_queue_action.assert_called_once()

    def test_supplement_chat_task_not_done_error(self, mock_task_lock):
        """Test supplementation fails when task is not done."""
        task_id = "test_task_123"
        supplement_data = SupplementChat(question="Add more details")
        mock_task_lock.status = Status.processing

        with patch(
            "app.controller.chat_controller.get_task_lock",
            return_value=mock_task_lock,
        ):
            with pytest.raises(UserException):
                supplement(task_id, supplement_data)

    @pytest.mark.asyncio
    async def test_stop_chat_success(self, mock_task_lock):
        """Test successful chat stopping."""
        task_id = "test_task_123"

        with (
            patch(
                "app.controller.chat_controller.get_task_lock_if_exists",
                return_value=mock_task_lock,
            ),
        ):
            response = await stop(task_id)

            assert isinstance(response, Response)
            assert response.status_code == 204
            mock_task_lock.put_queue.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_human_reply_success(self, mock_task_lock):
        """Test successful human reply."""
        task_id = "test_task_123"
        reply_data = HumanReply(agent="test_agent", reply="This is my reply")

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
        ):
            response = await human_reply(task_id, reply_data)

            assert isinstance(response, Response)
            assert response.status_code == 201
            mock_task_lock.put_human_input.assert_awaited_once_with(
                "test_agent", "This is my reply"
            )

    def test_install_mcp_success(self, mock_task_lock):
        """Test successful MCP installation."""
        task_id = "test_task_123"
        mcp_data: McpServers = {
            "mcpServers": {"test_server": {"config": "test"}}
        }

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
            patch(
                "app.controller.chat_controller._queue_action_from_worker"
            ) as mock_queue_action,
        ):
            response = install_mcp(task_id, mcp_data)

            assert isinstance(response, Response)
            assert response.status_code == 201
            mock_queue_action.assert_called_once()


@pytest.mark.integration
class TestChatControllerIntegration:
    """Integration tests for chat controller."""

    def test_chat_endpoint_integration(
        self, client: TestClient, sample_chat_data
    ):
        """Test chat endpoint through FastAPI test client."""
        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock"
            ) as mock_create_lock,
            patch(
                "app.controller.chat_controller.step_solve"
            ) as mock_step_solve,
            patch("app.controller.chat_controller.load_dotenv"),
            patch("app.controller.chat_controller.set_current_task_id"),
            patch(
                "app.controller.chat_controller._prepare_browser_for_request_with_timeout",
                new=AsyncMock(return_value=True),
            ),
            patch("pathlib.Path.mkdir"),
            patch("pathlib.Path.home", return_value=MagicMock()),
        ):
            mock_task_lock = MagicMock()
            mock_task_lock.put_queue = AsyncMock()
            mock_create_lock.return_value = mock_task_lock

            async def mock_generator():
                yield "data: test_response\n\n"

            mock_step_solve.return_value = mock_generator()

            response = client.post("/chat", json=sample_chat_data)

            assert response.status_code == 200
            assert (
                response.headers["content-type"]
                == "text/event-stream; charset=utf-8"
            )

    def test_improve_chat_endpoint_integration(self, client: TestClient):
        """Test improve chat endpoint through FastAPI test client."""
        task_id = "test_task_123"
        supplement_data = {"question": "Improve this code"}

        with (
            patch(
                "app.controller.chat_controller.get_task_lock"
            ) as mock_get_lock,
            patch(
                "app.controller.chat_controller._prepare_browser_for_request_with_timeout",
                new=AsyncMock(return_value=True),
            ),
        ):
            mock_task_lock = MagicMock()
            mock_task_lock.status = Status.processing
            mock_task_lock.put_queue = AsyncMock()
            mock_get_lock.return_value = mock_task_lock

            response = client.post(f"/chat/{task_id}", json=supplement_data)

            assert response.status_code == 201

    def test_supplement_chat_endpoint_integration(self, client: TestClient):
        """Test supplement chat endpoint through FastAPI test client."""
        task_id = "test_task_123"
        supplement_data = {"question": "Add more details"}

        with (
            patch(
                "app.controller.chat_controller.get_task_lock"
            ) as mock_get_lock,
            patch("app.controller.chat_controller._queue_action_from_worker"),
        ):
            mock_task_lock = MagicMock()
            mock_task_lock.status = Status.done
            mock_get_lock.return_value = mock_task_lock

            response = client.put(f"/chat/{task_id}", json=supplement_data)

            assert response.status_code == 201

    def test_stop_chat_endpoint_integration(self, client: TestClient):
        """Test stop chat endpoint through FastAPI test client."""
        task_id = "test_task_123"

        with (
            patch(
                "app.controller.chat_controller.get_task_lock_if_exists"
            ) as mock_get_lock,
        ):
            mock_task_lock = MagicMock()
            mock_task_lock.put_queue = AsyncMock()
            mock_get_lock.return_value = mock_task_lock

            response = client.delete(f"/chat/{task_id}")

            assert response.status_code == 204

    def test_human_reply_endpoint_integration(self, client: TestClient):
        """Test human reply endpoint through FastAPI test client."""
        task_id = "test_task_123"
        reply_data = {"agent": "test_agent", "reply": "This is my reply"}

        with (
            patch(
                "app.controller.chat_controller.get_task_lock"
            ) as mock_get_lock,
        ):
            mock_task_lock = MagicMock()
            mock_task_lock.put_human_input = AsyncMock()
            mock_get_lock.return_value = mock_task_lock

            response = client.post(
                f"/chat/{task_id}/human-reply", json=reply_data
            )

            assert response.status_code == 201

    def test_install_mcp_endpoint_integration(self, client: TestClient):
        """Test install MCP endpoint through FastAPI test client."""
        task_id = "test_task_123"
        mcp_data = {"mcpServers": {"test_server": {"config": "test"}}}

        with (
            patch(
                "app.controller.chat_controller.get_task_lock"
            ) as mock_get_lock,
            patch("app.controller.chat_controller._queue_action_from_worker"),
        ):
            mock_task_lock = MagicMock()
            mock_get_lock.return_value = mock_task_lock

            response = client.post(
                f"/chat/{task_id}/install-mcp", json=mcp_data
            )

            assert response.status_code == 201


@pytest.mark.model_backend
class TestChatControllerWithLLM:
    """Tests that require LLM backend (marked for selective running)."""

    @pytest.mark.asyncio
    async def test_post_with_real_llm_model(
        self, sample_chat_data, mock_request
    ):
        """Test chat endpoint with real LLM model (slow test)."""
        # This test would use actual LLM models and should be marked accordingly
        Chat(**sample_chat_data)

        # Test implementation would involve real model calls
        # This is marked as model_backend test for selective execution
        assert True  # Placeholder

    @pytest.mark.very_slow
    async def test_full_chat_workflow_with_llm(
        self, sample_chat_data, mock_request
    ):
        """Test complete chat workflow with LLM (very slow test)."""
        # This test would run the complete workflow including actual agent interactions
        # Marked as very_slow for execution only in full test mode
        assert True  # Placeholder


@pytest.mark.unit
class TestChatControllerErrorCases:
    """Test error cases and edge conditions."""

    @pytest.mark.asyncio
    async def test_post_with_invalid_data(self, mock_request):
        """Test chat endpoint with invalid data."""
        # Construction itself should raise a validation error due to multiple invalid fields
        with pytest.raises((ValueError, TypeError, ValidationError)):
            Chat(
                task_id="",  # Invalid empty task_id
                email="invalid_email",  # Invalid email format
                question="",  # Empty question
                attaches=[],
                model="invalid_model",  # Field not defined in model -> triggers error
                model_platform="invalid_platform",
                api_key="",
                api_url="invalid_url",
                new_agents=[],
                env_path="nonexistent.env",
                browser_port=-1,  # Invalid port
                summary_prompt="",
            )
        # If future validation moves to endpoint level, keep logic placeholder below.
        # (Intentionally not calling post with invalid Chat object since creation fails.)

    @pytest.mark.asyncio
    async def test_improve_with_nonexistent_task(self):
        """Test improve endpoint with nonexistent task."""
        task_id = "nonexistent_task"
        supplement_data = SupplementChat(question="Improve this code")
        request = SimpleNamespace()

        with patch(
            "app.controller.chat_controller.get_task_lock",
            side_effect=KeyError("Task not found"),
        ):
            with pytest.raises(KeyError):
                await improve(task_id, supplement_data, request)

    def test_supplement_with_empty_question(self, mock_task_lock):
        """Test supplement endpoint with empty question."""
        task_id = "test_task_123"
        supplement_data = SupplementChat(question="")
        mock_task_lock.status = Status.done

        with (
            patch(
                "app.controller.chat_controller.get_task_lock",
                return_value=mock_task_lock,
            ),
            patch("app.controller.chat_controller._queue_action_from_worker"),
        ):
            # Should handle empty question gracefully or raise appropriate error
            response = supplement(task_id, supplement_data)
            assert response.status_code == 201  # Or should it be an error?

    @pytest.mark.asyncio
    async def test_post_environment_setup_failure(
        self, sample_chat_data, mock_request
    ):
        """Test chat endpoint when environment setup fails."""
        chat_data = Chat(**sample_chat_data)

        with (
            patch(
                "app.controller.chat_controller.get_or_create_task_lock"
            ) as mock_create_lock,
            patch(
                "app.controller.chat_controller.sanitize_env_path",
                return_value="/tmp/fake.env",
            ),
            patch(
                "app.controller.chat_controller.load_dotenv",
                side_effect=Exception("Env load failed"),
            ),
            patch(
                "pathlib.Path.mkdir",
                side_effect=Exception("Directory creation failed"),
            ),
        ):
            mock_task_lock = MagicMock()
            mock_create_lock.return_value = mock_task_lock

            # Should handle environment setup failures gracefully
            with pytest.raises(Exception):
                await post(chat_data, mock_request)