microsoft--magentic-ui/tests/test_start_mounted_folder_metadata.py

from __future__ import annotations

import json
from typing import Any
from unittest.mock import AsyncMock, MagicMock

import pytest

from magentic_ui.backend.datamodel import RunStatus
from magentic_ui.backend.web.managers.connection import WebSocketManager


async def _empty_stream(**_: Any):
    if False:
        yield None


@pytest.mark.asyncio
async def test_start_message_persists_mounted_folder_metadata() -> None:
    team = MagicMock()
    team.set_uploaded_file_infos = MagicMock()
    team.run_stream = MagicMock(side_effect=lambda **kwargs: _empty_stream(**kwargs))
    team.close = AsyncMock()

    run = MagicMock(id=1, status=RunStatus.CREATED)
    mgr = WebSocketManager.__new__(WebSocketManager)
    mgr._connections = {1: MagicMock()}
    mgr._closed_connections = set()
    mgr._quicksand_manager = None
    mgr._team_managers = {1: team}
    mgr._stream_tasks = {}
    mgr._stopping_runs = set()
    mgr._get_run = AsyncMock(return_value=run)
    mgr._update_run_status = AsyncMock()
    mgr._send_message = AsyncMock()
    mgr._save_message = AsyncMock()
    mgr._run_is_mid_input = AsyncMock(return_value=False)
    mgr.db_manager = MagicMock()

    folder_path = "/home/user/project"

    await mgr._execute_stream(
        1,
        "Inspect this repo",
        settings_config={"mount_dirs": [folder_path]},
    )

    broadcast_payload = mgr._send_message.await_args_list[0].args[1]
    metadata = broadcast_payload["data"]["metadata"]
    assert json.loads(metadata["attached_files"]) == []
    assert metadata["mounted_folder"] == {
        "name": "project",
        "path": folder_path,
    }

    save_payload = mgr._save_message.await_args_list[0].args[1]
    assert save_payload["metadata"]["mounted_folder"] == {
        "name": "project",
        "path": folder_path,
    }

    team.run_stream.assert_called_once()
    assert team.run_stream.call_args.kwargs["mount_dirs"] == [folder_path]