test_cli.py

import json
import os
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from io import StringIO
from pathlib import Path
from textwrap import dedent
from typing import Any, TypeVar
from unittest import mock
from unittest.mock import Mock, patch

import pytest
import responses
import yaml
from bluesky.protocols import Movable
from bluesky_stomp.messaging import StompClient
from bluesky_stomp.models import MessageTopic
from click.testing import CliRunner
from ophyd_async.core import AsyncStatus
from pydantic import BaseModel
from requests.exceptions import ConnectionError
from responses import matchers
from stomp.connect import StompConnection11 as Connection

from blueapi import __version__
from blueapi.cli.cli import ParametersType, main
from blueapi.cli.format import OutputFormat, fmt_dict
from blueapi.client.event_bus import BlueskyStreamingError
from blueapi.client.rest import (
    BlueskyRemoteControlError,
    InvalidParametersError,
    ParameterError,
    UnauthorisedAccessError,
    UnknownPlanError,
)
from blueapi.config import (
    ApplicationConfig,
    ScratchConfig,
    ScratchRepository,
    generate_config_schema,
)
from blueapi.core.bluesky_types import DataEvent, Plan
from blueapi.service.model import (
    DeviceModel,
    DeviceResponse,
    EnvironmentResponse,
    PlanModel,
    PlanResponse,
    PythonEnvironmentResponse,
    TaskRequest,
    TaskResponse,
)
from blueapi.worker.event import (
    ProgressEvent,
    TaskError,
    TaskResult,
    TaskStatus,
    WorkerEvent,
    WorkerState,
)


@pytest.fixture
def mock_connection() -> Mock:
    return Mock(spec=Connection)


@pytest.fixture
def mock_stomp_client(mock_connection: Mock) -> StompClient:
    return StompClient(conn=mock_connection)


@pytest.fixture
def runner():
    return CliRunner()


def test_cli_version(runner: CliRunner):
    result = runner.invoke(main, ["--version"])
    assert result.stdout == f"blueapi, version {__version__}\n"


def test_main_no_params():
    runner = CliRunner()
    result = runner.invoke(main)
    expected = "Please invoke subcommand!\n"

    assert result.stdout == expected


@patch("blueapi.service.main.start")
@patch("blueapi.cli.scratch.setup_scratch")
@patch("blueapi.cli.cli.os.umask")
@pytest.mark.parametrize("subcommand", ["serve", "setup-scratch"])
def test_runs_with_umask_002(
    mock_umask: Mock,
    mock_setup_scratch: Mock,
    mock_start: Mock,
    runner: CliRunner,
    subcommand: str,
):
    runner.invoke(main, [subcommand])
    mock_umask.assert_called_once_with(0o002)


@patch("blueapi.client.rest.requests.Session.request")
def test_connection_error_caught_by_wrapper_func(
    mock_requests: Mock, runner: CliRunner
):
    mock_requests.side_effect = ConnectionError()
    result = runner.invoke(main, ["controller", "plans"])

    assert result.output == "Error: Failed to establish connection to blueapi server.\n"


@patch("blueapi.client.rest.requests.Session.request")
def test_authentication_error_caught_by_wrapper_func(
    mock_requests: Mock, runner: CliRunner
):
    mock_requests.side_effect = UnauthorisedAccessError(message="<Response [401]>")
    result = runner.invoke(main, ["controller", "plans"])
    assert (
        result.output
        == "Error: Access denied. Please check your login status and try again.\n"
    )


@patch("blueapi.client.rest.requests.Session.request")
def test_remote_error_raised_by_wrapper_func(mock_requests: Mock, runner: CliRunner):
    mock_requests.side_effect = BlueskyRemoteControlError("Response [450]")
    result = runner.invoke(main, ["controller", "plans"])
    assert (
        isinstance(result.exception, BlueskyRemoteControlError)
        and result.exception.args == ("Response [450]",)
        and result.exit_code == 1
    )


class MyModel(BaseModel):
    id: str


ComplexType = TypeVar("ComplexType")


@dataclass
class MyDevice(Movable[ComplexType]):
    name: str

    @AsyncStatus.wrap
    async def set(self, value: ComplexType): ...


@responses.activate
def test_get_plans(runner: CliRunner):
    plan = Plan(name="my-plan", model=MyModel)

    response = responses.add(
        responses.GET,
        "http://localhost:8000/plans",
        json=PlanResponse(plans=[PlanModel.from_plan(plan)]).model_dump(),
        status=200,
    )

    plans = runner.invoke(main, ["controller", "plans"])
    assert response.call_count == 1
    assert plans.output == "my-plan\n    Args\n      id=string (Required)\n"


@responses.activate
def test_get_devices(runner: CliRunner):
    device = MyDevice(name="my-device")

    response = responses.add(
        responses.GET,
        "http://localhost:8000/devices",
        json=DeviceResponse(devices=[DeviceModel.from_device(device)]).model_dump(),
        status=200,
    )

    plans = runner.invoke(main, ["controller", "devices"])
    assert response.call_count == 1
    assert plans.output == "my-device\n    Movable['ComplexType']\n"


def test_invalid_config_path_handling(runner: CliRunner):
    # test what happens if you pass an invalid config file...
    result = runner.invoke(main, ["-c", "non_existent.yaml"])
    assert result.exit_code == 1


@patch("blueapi.cli.cli.BlueapiClient.plans")
@patch("blueapi.cli.cli.OutputFormat.FULL.display")
def test_options_via_env(mock_display, mock_plans, runner: CliRunner):
    result = runner.invoke(
        main, args=["controller", "plans"], env={"BLUEAPI_CONTROLLER_OUTPUT": "full"}
    )

    mock_plans.__iter__.assert_called_once_with()
    mock_display.assert_called_once_with(PlanResponse(plans=list(mock_plans)))
    assert result.exit_code == 0


def test_invalid_config_via_env(runner: CliRunner):
    result = runner.invoke(main, env={"BLUEAPI_CONFIG": "non_existent.yaml"})
    assert result.exit_code == 1


@responses.activate
def test_submit_plan(runner: CliRunner):
    body_data = {
        "name": "sleep",
        "params": {"time": 5},
        "instrument_session": "cm12345-1",
    }

    response = responses.post(
        url="http://a.fake.host:12345/tasks",
        match=[matchers.json_params_matcher(body_data)],
    )

    config_path = "tests/unit_tests/example_yaml/rest_and_stomp_config.yaml"
    output = runner.invoke(
        main,
        [
            "-c",
            config_path,
            "controller",
            "run",
            "-i",
            "cm12345-1",
            "sleep",
            '{"time": 5}',
        ],
    )

    assert response.call_count == 1, output.output


@responses.activate
def test_submit_plan_without_stomp(runner: CliRunner):
    config_path = "tests/unit_tests/example_yaml/rest_config.yaml"
    result = runner.invoke(
        main,
        [
            "-c",
            config_path,
            "controller",
            "run",
            "-i",
            "cm12345-1",
            "sleep",
            '{"time": 5}',
        ],
    )

    assert (
        result.stderr
        == "Error: Stomp configuration required to run plans is missing or disabled\n"
    )


@patch("blueapi.client.client.StompClient")
@responses.activate
def test_run_plan(stomp_client: StompClient, runner: CliRunner):
    task_id = "abcd-1234"
    submit_response = responses.post(
        url="http://a.fake.host:12345/tasks",
        match=[
            matchers.json_params_matcher(
                {
                    "name": "sleep",
                    "params": {"time": 3},
                    "instrument_session": "cm12345-1",
                }
            )
        ],
        json={"task_id": task_id},
        status=201,
    )
    run_response = responses.put(
        url="http://a.fake.host:12345/worker/task",
        match=[matchers.json_params_matcher({"task_id": task_id})],
        json={"task_id": task_id},
    )

    def mock_events(topic: MessageTopic, callback: Callable[[Any, Any], Any]):
        if topic.name != "public.worker.event":
            return
        ctx = Mock()
        ctx.correlation_id = task_id
        callback(
            WorkerEvent(
                state=WorkerState.RUNNING,
                task_status=TaskStatus(
                    task_id=task_id,
                    task_complete=False,
                    task_failed=False,
                    result=None,
                ),
            ),
            ctx,
        )
        callback(ProgressEvent(task_id=task_id), ctx)
        callback(DataEvent(name="event", doc={}, task_id=task_id), ctx)
        callback(
            WorkerEvent(
                state=WorkerState.IDLE,
                task_status=TaskStatus(
                    task_id=task_id, task_complete=False, task_failed=False, result=None
                ),
            ),
            ctx,
        )
        callback(
            WorkerEvent(
                state=WorkerState.IDLE,
                task_status=TaskStatus(
                    task_id=task_id,
                    task_complete=True,
                    task_failed=False,
                    result=TaskResult(result=None, type="NoneType"),
                ),
            ),
            ctx,
        )

    stomp = stomp_client.for_broker(...)  # type: ignore
    stomp.subscribe.side_effect = mock_events  # type: ignore

    config_path = "tests/unit_tests/example_yaml/rest_and_stomp_config.yaml"
    result = runner.invoke(
        main,
        [
            "-c",
            config_path,
            "controller",
            "run",
            "-i",
            "cm12345-1",
            "sleep",
            '{"time": 3}',
        ],
    )

    assert result.exit_code == 0
    assert submit_response.call_count == 1
    assert run_response.call_count == 1


@pytest.mark.parametrize(
    "result,failed,message",
    [
        (TaskResult(result=None, type="NoneType"), False, "Plan succeeded\n"),
        (TaskResult(result=32, type="int"), False, "Plan succeeded: 32\n"),
        (
            TaskResult(result=None, type="CustomType"),
            False,
            "Plan returned unserializable result of type 'CustomType'\n",
        ),
        (
            TaskError(type="ValueError", message="Error with value"),
            True,
            "Plan failed: ValueError: Error with value\n",
        ),
    ],
)
@patch("blueapi.cli.cli.BlueapiClient")
def test_run_plan_feedback(
    mock_client: Mock,
    runner: CliRunner,
    result: TaskResult | TaskError | None,
    failed: bool,
    message: str,
):
    bc = mock_client.from_config()
    bc.run_task.return_value = TaskStatus(
        task_id="foo_bar",
        task_complete=True,
        task_failed=failed,
        result=result,
    )
    res = runner.invoke(
        main,
        ["controller", "run", "-i", "cm12345-1", "name"],
    )
    bc.run_task.assert_called_once_with(
        TaskRequest(name="name", params={}, instrument_session="cm12345-1"),
        on_event=mock.ANY,
    )
    assert res.exit_code == 0
    assert res.stdout == message


@responses.activate
def test_run_plan_background_without_stomp(runner: CliRunner):
    submit_response = responses.post(
        url="http://a.fake.host:12345/tasks",
        match=[
            matchers.json_params_matcher(
                {
                    "name": "sleep",
                    "params": {"time": 3},
                    "instrument_session": "cm12345-1",
                }
            )
        ],
        json={"task_id": "abcd-1234"},
        status=201,
    )
    run_response = responses.put(
        url="http://a.fake.host:12345/worker/task",
        match=[matchers.json_params_matcher({"task_id": "abcd-1234"})],
        json={"task_id": "abcd-1234"},
    )

    config_path = "tests/unit_tests/example_yaml/rest_config.yaml"
    result = runner.invoke(
        main,
        [
            "-c",
            config_path,
            "controller",
            "run",
            "-i",
            "cm12345-1",
            "--background",
            "sleep",
            '{"time": 3}',
        ],
    )

    assert result.exit_code == 0
    assert result.output == "abcd-1234\n"
    assert submit_response.call_count == 1
    assert run_response.call_count == 1


def test_invalid_stomp_config_for_listener(runner: CliRunner):
    result = runner.invoke(main, ["controller", "listen"])
    assert isinstance(result.exception, BlueskyStreamingError)
    assert str(result.exception) == "Message bus needs to be configured"


def test_cannot_run_plans_without_stomp_config(runner: CliRunner):
    result = runner.invoke(
        main,
        [
            "controller",
            "run",
            "-i",
            "cm12345-1",
            "sleep",
            '{"time": 5}',
        ],
    )
    assert result.exit_code == 1
    assert (
        result.stderr
        == "Error: Stomp configuration required to run plans is missing or disabled\n"
    )


def test_cannot_start_a_plan_without_an_instrument_session(runner: CliRunner):
    result = runner.invoke(
        main,
        [
            "controller",
            "run",
            "--background",
            "sleep",
            '{"time": 5}',
        ],
    )
    assert result.exit_code == 2
    assert "Error: Missing option '-i' / '--instrument-session'.\n" in result.stderr


@patch("blueapi.client.rest.BlueapiRestClient.update_worker_task")
@patch("blueapi.client.rest.BlueapiRestClient.create_task")
def test_can_pass_an_instrument_session_with_an_environment_variable(
    mock_create_task: Mock, mock_update_worker_task: Mock, runner: CliRunner
):
    mock_create_task.return_value = TaskResponse(task_id="foo")
    mock_update_worker_task.return_value = TaskResponse(task_id="foo")
    with patch.dict(
        os.environ,
        {"BLUEAPI_CONTROLLER_RUN_INSTRUMENT_SESSION": "cm12345-1"},
        clear=True,
    ):
        # assert visit passed to rest
        result = runner.invoke(
            main,
            [
                "controller",
                "run",
                "--background",
                "sleep",
                '{"time": 5.0}',
            ],
        )
    assert result.exit_code == 0, result.output
    mock_create_task.assert_called_once_with(
        TaskRequest(
            name="sleep",
            params={"time": 5.0},
            instrument_session="cm12345-1",
        )
    )


@patch("blueapi.cli.cli.StompClient")
def test_valid_stomp_config_for_listener(
    mock_stomp_client: StompClient,
    runner: CliRunner,
    mock_connection: Mock,
):
    mock_connection.is_connected.return_value = True
    result = runner.invoke(
        main,
        [
            "-c",
            "tests/unit_tests/example_yaml/valid_stomp_config.yaml",
            "controller",
            "listen",
        ],
        input="\n",
    )
    assert result.output == dedent("""\
                Subscribing to all bluesky events from localhost:61613
                Press enter to exit
                """)
    assert result.exit_code == 0


@responses.activate
def test_get_env(runner: CliRunner):
    environment_id = uuid.uuid4()
    responses.add(
        responses.GET,
        "http://localhost:8000/environment",
        json=EnvironmentResponse(
            environment_id=environment_id, initialized=True
        ).model_dump(mode="json"),
        status=200,
    )

    env = runner.invoke(main, ["controller", "env"])
    assert (
        env.output == f"environment_id=UUID('{environment_id}') "
        "initialized=True "
        "error_message=None\n"
    )


@responses.activate
def test_get_state(runner: CliRunner):
    responses.add(
        responses.GET, "http://localhost:8000/worker/state", json="IDLE", status=200
    )
    state = runner.invoke(main, ["controller", "state"])
    print(state.stderr)
    assert state.exit_code == 0
    assert state.output == "IDLE\n"


@responses.activate(assert_all_requests_are_fired=True)
@patch("blueapi.client.client.time.sleep", return_value=None)
def test_reset_env_client_behavior(
    mock_sleep: Mock,
    runner: CliRunner,
):
    environment_id = uuid.uuid4()
    responses.add(
        responses.DELETE,
        "http://localhost:8000/environment",
        json=EnvironmentResponse(
            environment_id=environment_id, initialized=False
        ).model_dump(mode="json"),
        status=200,
    )

    env_state = [False, False, True]
    environment_id = uuid.uuid4()
    for state in env_state:
        responses.add(
            responses.GET,
            "http://localhost:8000/environment",
            json=EnvironmentResponse(
                environment_id=environment_id, initialized=state
            ).model_dump(mode="json"),
            status=200,
        )

    # Invoke the CLI command that would trigger the environment initialization check
    reload_result = runner.invoke(main, ["controller", "env", "-r"])

    # Verify if sleep was called between polling iterations
    mock_sleep.assert_called()

    for index, call in enumerate(responses.calls):
        if index == 0:
            assert call.request.method == "DELETE"
            assert call.request.url == "http://localhost:8000/environment"
        else:
            assert call.request.method == "GET"
            assert call.request.url == "http://localhost:8000/environment"

    # Check if the final environment status is printed correctly
    # assert "Environment is initialized." in result.output
    assert reload_result.output == dedent(f"""\
                Reloading environment
                Environment is initialized
                environment_id=UUID('{environment_id}') initialized=True error_message=None
                """)  # noqa: E501


@responses.activate
@patch("blueapi.client.client.time.sleep", return_value=None)
def test_env_timeout(mock_sleep: Mock, runner: CliRunner):
    # Setup mocked responses for the REST endpoints
    environment_id = uuid.uuid4()
    responses.add(
        responses.DELETE,
        "http://localhost:8000/environment",
        status=200,
        json=EnvironmentResponse(
            environment_id=environment_id, initialized=False
        ).model_dump(mode="json"),
    )
    # Add responses for each polling attempt, all indicating not initialized
    responses.add(
        responses.GET,
        "http://localhost:8000/environment",
        json=EnvironmentResponse(
            environment_id=environment_id, initialized=False
        ).model_dump(mode="json"),
        status=200,
    )

    # Run the command that should interact with these endpoints
    result = runner.invoke(main, ["controller", "env", "-r", "-t", "0.1"])
    if result.exception is not None:
        assert isinstance(result.exception, TimeoutError), "Expected a TimeoutError"
        assert (
            result.exception.args[0]
            == "Failed to reload the environment within 0.1 seconds, "
            "a server restart is recommended"
        )
    else:
        raise AssertionError("Expected an exception but got None")

    # First call should be DELETE
    assert responses.calls[0].request.method == "DELETE"
    assert responses.calls[0].request.url == "http://localhost:8000/environment"

    # Remaining calls should all be GET
    for call in responses.calls[1:]:  # Skip the first DELETE request # type: ignore
        assert call.request.method == "GET"
        assert call.request.url == "http://localhost:8000/environment"

    # Check the output for the timeout message
    assert result.output == "Reloading environment\n"
    assert (
        result.exit_code == 1
    )  # Assuming your command exits successfully even on timeout for simplicity


@responses.activate
def test_env_reload_server_side_error(runner: CliRunner):
    # Setup mocked error response from the server
    responses.add(
        responses.DELETE, "http://localhost:8000/environment", status=500, json={}
    )

    result = runner.invoke(main, ["controller", "env", "-r"])
    assert isinstance(result.exception, BlueskyRemoteControlError), (
        "Expected a BlueskyRemoteError from cli runner"
    )
    assert result.exception.args[0] == "Failed to tear down the environment"

    # Check if the endpoints were hit as expected
    assert len(responses.calls) == 1  # +1 for the DELETE call

    # Only call should be DELETE
    assert responses.calls[0].request.method == "DELETE"
    assert responses.calls[0].request.url == "http://localhost:8000/environment"

    # Check the output for the timeout message
    # TODO this seems wrong but this is the current behaviour
    # There should be an error message
    assert result.output == "Reloading environment\n"

    assert result.exit_code == 1


@pytest.mark.parametrize(
    "exception, error_message",
    [
        (UnknownPlanError(), "Error: Plan 'sleep' was not recognised\n"),
        (UnauthorisedAccessError(), "Error: Unauthorised request\n"),
        (
            InvalidParametersError(
                errors=[
                    ParameterError(
                        loc=["body", "params", "foo"],
                        type="missing",
                        msg="Foo is missing",
                        input=None,
                    )
                ]
            ),
            "Error: Incorrect parameters supplied\n    Missing value for 'foo'\n",
        ),
        (
            BlueskyRemoteControlError("Server error"),
            "Error: remote control error: Server error\n",
        ),
        (
            ValueError("Error parsing parameters"),
            "Error: task could not run: Error parsing parameters\n",
        ),
        (
            BlueskyStreamingError("streaming failed"),
            "Error: streaming error: streaming failed\n",
        ),
    ],
    ids=[
        "unknown_plan",
        "unauthorised_access",
        "invalid_parameters",
        "remote_control",
        "value_error",
        "streaming_error",
    ],
)
def test_error_handling(exception, error_message, runner: CliRunner):
    # Patching the create_task method to raise different exceptions
    with patch(
        "blueapi.client.rest.BlueapiRestClient.create_task", side_effect=exception
    ):
        result = runner.invoke(
            main,
            [
                "-c",
                "tests/unit_tests/example_yaml/valid_stomp_config.yaml",
                "controller",
                "run",
                "-i",
                "cm12345-1",
                "sleep",
                '{"time": 5}',
            ],
        )
    assert result.stderr == error_message
    assert result.exit_code == 1


@pytest.mark.parametrize(
    "params, error",
    [
        ("{", "Parameters are not valid JSON"),
        ("[]", "Parameters must be a JSON object with string keys"),
    ],
)
def test_run_task_parsing_errors(params: str, error: str, runner: CliRunner):
    result = runner.invoke(
        main,
        [
            "-c",
            "tests/unit_tests/example_yaml/valid_stomp_config.yaml",
            "controller",
            "run",
            "-i",
            "cm12345-1",
            "sleep",
            params,
        ],
    )
    assert error in result.stderr
    assert result.exit_code == 2


def test_device_output_formatting():
    """Test for alternative device output formats"""

    device = MyDevice("my-device")

    devices = DeviceResponse(devices=[DeviceModel.from_device(device)])

    compact = dedent("""\
                my-device
                    Movable['ComplexType']
                """)

    _assert_matching_formatting(OutputFormat.COMPACT, devices, compact)

    json_out = dedent("""\
                [
                  {
                    "name": "my-device",
                    "protocols": [
                      {
                        "name": "Movable",
                        "types": [
                          "ComplexType"
                        ]
                      }
                    ]
                  }
                ]
                """)
    _assert_matching_formatting(OutputFormat.JSON, devices, json_out)
    _ = json.loads(json_out)

    full = dedent("""\
            my-device
                Movable['ComplexType']
            """)
    _assert_matching_formatting(OutputFormat.FULL, devices, full)


class ExtendedModel(BaseModel):
    name: str
    keys: list[int]
    metadata: None | Mapping[str, str] = None


def test_plan_output_formatting():
    """Test for alternative plan output formats"""

    plan = Plan(
        name="my-plan",
        description=dedent("""\
            Summary of description

            Rest of description
            """),
        model=ExtendedModel,
    )
    plans = PlanResponse(plans=[PlanModel.from_plan(plan)])

    compact = dedent("""\
                my-plan
                    Summary of description
                    Args
                      name=string (Required)
                      keys=[integer] (Required)
                      metadata=object
                """)

    _assert_matching_formatting(OutputFormat.COMPACT, plans, compact)

    json_out = dedent("""\
            [
              {
                "name": "my-plan",
                "description": "Summary of description\\n\\nRest of description\\n",
                "parameter_schema": {
                  "properties": {
                    "name": {
                      "title": "Name",
                      "type": "string"
                    },
                    "keys": {
                      "items": {
                        "type": "integer"
                      },
                      "title": "Keys",
                      "type": "array"
                    },
                    "metadata": {
                      "anyOf": [
                        {
                          "additionalProperties": {
                            "type": "string"
                          },
                          "type": "object"
                        },
                        {
                          "type": "null"
                        }
                      ],
                      "default": null,
                      "title": "Metadata"
                    }
                  },
                  "required": [
                    "name",
                    "keys"
                  ],
                  "title": "ExtendedModel",
                  "type": "object"
                }
              }
            ]
                """)
    _assert_matching_formatting(OutputFormat.JSON, plans, json_out)
    _ = json.loads(json_out)

    full = dedent("""\
        my-plan
            Summary of description

            Rest of description
            Schema
                {
                  "properties": {
                    "name": {
                      "title": "Name",
                      "type": "string"
                    },
                    "keys": {
                      "items": {
                        "type": "integer"
                      },
                      "title": "Keys",
                      "type": "array"
                    },
                    "metadata": {
                      "anyOf": [
                        {
                          "additionalProperties": {
                            "type": "string"
                          },
                          "type": "object"
                        },
                        {
                          "type": "null"
                        }
                      ],
                      "default": null,
                      "title": "Metadata"
                    }
                  },
                  "required": [
                    "name",
                    "keys"
                  ],
                  "title": "ExtendedModel",
                  "type": "object"
                }
            """)
    _assert_matching_formatting(OutputFormat.FULL, plans, full)


def test_event_formatting():
    data = DataEvent(
        name="start",
        doc={"foo": "bar", "fizz": {"buzz": (1, 2, 3), "hello": "world"}},
        task_id="0000-1111",
    )
    worker = WorkerEvent(
        state=WorkerState.RUNNING,
        task_status=TaskStatus(
            task_id="count", task_complete=False, task_failed=False, result=None
        ),
        errors=[],
        warnings=[],
    )
    progress = ProgressEvent(task_id="start", statuses={})

    _assert_matching_formatting(
        OutputFormat.JSON,
        data,
        (
            """{"name": "start", "doc": """
            """{"foo": "bar", "fizz": {"buzz": [1, 2, 3], "hello": "world"}}, """
            """"task_id": "0000-1111"}\n"""
        ),
    )
    _assert_matching_formatting(OutputFormat.COMPACT, data, "Data Event: start\n")
    _assert_matching_formatting(
        OutputFormat.FULL,
        data,
        dedent("""\
            Start:
                foo: bar
                fizz:
                    buzz: (1, 2, 3)
                    hello: world
            """),
    )

    _assert_matching_formatting(
        OutputFormat.JSON,
        worker,
        (
            '{"state": "RUNNING", "task_status": {'
            '"task_id": "count", '
            '"result": null, '
            '"task_complete": false, '
            '"task_failed": false'
            '}, "errors": [], "warnings": []}\n'
        ),
    )
    _assert_matching_formatting(OutputFormat.COMPACT, worker, "Worker Event: RUNNING\n")
    _assert_matching_formatting(
        OutputFormat.FULL,
        worker,
        "WorkerEvent: RUNNING\n    task_id: count\n",
    )

    _assert_matching_formatting(
        OutputFormat.JSON, progress, """{"task_id": "start", "statuses": {}}\n"""
    )
    _assert_matching_formatting(OutputFormat.COMPACT, progress, "Progress: ???%\n")
    _assert_matching_formatting(
        OutputFormat.FULL, progress, "Progress:\n    task_id: start\n"
    )


def test_unknown_object_formatting():
    demo = {"foo": 42, "bar": ["hello", "World"]}

    exp = """{"foo": 42, "bar": ["hello", "World"]}\n"""
    _assert_matching_formatting(OutputFormat.JSON, demo, exp)

    exp = """{'bar': ['hello', 'World'], 'foo': 42}\n"""
    _assert_matching_formatting(OutputFormat.COMPACT, demo, exp)

    _assert_matching_formatting(OutputFormat.FULL, demo, exp)


def test_dict_formatting():
    demo = {"name": "foo", "keys": [1, 2, 3], "metadata": {"fizz": "buzz"}}
    exp = """\nname: foo\nkeys: [1, 2, 3]\nmetadata:\n    fizz: buzz"""
    assert fmt_dict(demo, 0) == exp

    demo = "not a dict"
    assert fmt_dict(demo, 0) == " not a dict"


def test_generic_base_model_formatting():
    model = ExtendedModel(name="demo", keys=[1, 2, 3], metadata={"fizz": "buzz"})
    exp = '{"name": "demo", "keys": [1, 2, 3], "metadata": {"fizz": "buzz"}}\n'
    _assert_matching_formatting(OutputFormat.JSON, model, exp)

    _assert_matching_formatting(
        OutputFormat.FULL,
        model,
        dedent("""\
            ExtendedModel
                name: demo
                keys: [1, 2, 3]
                metadata:
                    fizz: buzz
            """),
    )


@patch("blueapi.cli.cli.setup_scratch")
def test_init_scratch_calls_setup_scratch(mock_setup_scratch: Mock, runner: CliRunner):
    expected_config = ScratchConfig(
        root=Path("/tmp"),
        repositories=[
            ScratchRepository(
                name="dodal",
                remote_url="https://github.com/DiamondLightSource/dodal.git",
            )
        ],
    )

    result = runner.invoke(
        main,
        ["-c", "tests/unit_tests/example_yaml/scratch.yaml", "setup-scratch"],
        input="\n",
    )
    assert result.exit_code == 0
    mock_setup_scratch.assert_called_once_with(expected_config)


def _assert_matching_formatting(fmt: OutputFormat, obj: Any, expected: str):
    output = StringIO()
    fmt.display(obj, output)
    assert expected == output.getvalue()


def test_login_success(
    runner: CliRunner,
    config_with_auth: str,
    mock_authn_server: responses.RequestsMock,
):
    with patch("webbrowser.open_new_tab", return_value=False):
        result = runner.invoke(main, ["-c", config_with_auth, "login"])
        assert (
            "Logging in\n"
            "Please login from this URL:- https://example.com/verify\n"
            "Logged in and cached new token\n" == result.output
        )
        assert result.exit_code == 0


def test_token_login_with_valid_token(
    runner: CliRunner,
    config_with_auth: str,
    mock_authn_server: responses.RequestsMock,
    cached_valid_token: Path,
):
    result = runner.invoke(main, ["-c", config_with_auth, "login"])
    assert "Logged in\n" == result.output
    assert result.exit_code == 0


def test_login_with_refresh_token(
    runner: CliRunner,
    config_with_auth: str,
    mock_authn_server: responses.RequestsMock,
    cached_valid_refresh: Path,
):
    result = runner.invoke(main, ["-c", config_with_auth, "login"])

    assert "Logged in\n" == result.output
    assert result.exit_code == 0


def test_login_when_cached_token_decode_fails(
    runner: CliRunner,
    config_with_auth: str,
    mock_authn_server: responses.RequestsMock,
    cached_expired_refresh: Path,
):
    with patch("webbrowser.open_new_tab", return_value=False):
        result = runner.invoke(main, ["-c", config_with_auth, "login"])
        assert (
            "Logging in\n"
            "Please login from this URL:- https://example.com/verify\n"
            "Logged in and cached new token\n" in result.output
        )
        assert result.exit_code == 0


def test_login_with_unauthenticated_server(
    runner: CliRunner,
    config_with_auth: str,
    mock_unauthenticated_server: responses.RequestsMock,
):
    result = runner.invoke(main, ["-c", config_with_auth, "login"])
    assert "Server is not configured to use authentication!\n" == result.output
    assert result.exit_code == 0


def test_logout_success(
    runner: CliRunner,
    config_with_auth: str,
    cached_valid_refresh: Path,
    mock_authn_server: responses.RequestsMock,
):
    assert cached_valid_refresh.exists()
    result = runner.invoke(main, ["-c", config_with_auth, "logout"])
    assert "Logged out" in result.output
    assert not cached_valid_refresh.exists()


def test_logout_invalid_token(runner: CliRunner):
    with patch("blueapi.cli.cli.SessionManager") as sm:
        sm.from_cache.side_effect = ValueError("Invalid token")
        result = runner.invoke(main, ["logout"])

    assert result.exit_code == 1
    assert (
        result.output
        == "Error: Login token is not valid - remove before trying again\n"
    )


def test_logout_unknown_error(runner: CliRunner):
    with patch("blueapi.cli.cli.SessionManager") as sm:
        sm.from_cache.side_effect = Exception("Invalid token")
        result = runner.invoke(main, ["logout"])

    assert result.exit_code == 1
    assert result.output == "Error: Error logging out: Invalid token\n"


def test_logout_when_no_cache(
    runner: CliRunner,
    config_with_auth: str,
):
    result = runner.invoke(main, ["-c", config_with_auth, "logout"])
    assert "Logged out" in result.output


def test_local_cache_cleared_on_logout_when_oidc_unavailable(
    runner: CliRunner,
    config_with_auth: str,
    cached_valid_refresh: Path,
):
    assert cached_valid_refresh.exists()
    result = runner.invoke(main, ["-c", config_with_auth, "logout"])
    assert (
        "An unexpected error occurred while attempting to log out from the server."
        in result.output
    )
    assert not cached_valid_refresh.exists()


def test_wrapper_is_a_directory_error(
    runner: CliRunner, mock_authn_server: responses.RequestsMock, tmp_path
):
    config: ApplicationConfig = ApplicationConfig(auth_token_path=tmp_path)
    config_path = tmp_path / "config.yaml"
    with open(config_path, mode="w") as valid_auth_config_file:
        valid_auth_config_file.write(yaml.dump(config.model_dump()))
    result = runner.invoke(main, ["-c", config_path.as_posix(), "login"])
    assert (
        "Invalid path: a directory path was provided instead of a file path\n"
        == result.stdout
    )


def test_wrapper_permission_error(
    runner: CliRunner, mock_authn_server: responses.RequestsMock, tmp_path
):
    token_file: Path = tmp_path / "dir/token"

    config: ApplicationConfig = ApplicationConfig(auth_token_path=token_file)
    config_path = tmp_path / "config.yaml"
    with open(config_path, mode="w") as valid_auth_config_file:
        valid_auth_config_file.write(yaml.dump(config.model_dump()))
    with patch.object(Path, "write_text", side_effect=PermissionError):
        result = runner.invoke(main, ["-c", config_path.as_posix(), "login"])
    assert f"Permission denied: Cannot write to {token_file}\n" == result.stdout


@responses.activate
def test_get_python_environment(runner: CliRunner):
    scratch_config = {
        "installed_packages": [
            {
                "name": "bar",
                "version": "0.0.1",
                "location": "/tmp/bar",
                "is_dirty": "false",
                "source": "pypi",
            },
            {
                "name": "foo",
                "version": "https://github.com/example/foo.git @18ec206e",
                "location": "/tmp/foo",
                "is_dirty": "true",
                "source": "scratch",
            },
        ],
        "scratch_enabled": "true",
    }
    response = responses.add(
        responses.GET,
        "http://localhost:8000/python_environment",
        json=scratch_config,
        status=200,
    )

    result = runner.invoke(main, ["controller", "get-python-env"])
    assert response.call_count == 1
    assert result.exit_code == 0

    assert result.output == dedent("""\
        Scratch Status: enabled
        - bar @ (0.0.1)
        - foo @ (https://github.com/example/foo.git @18ec206e) (Dirty) (Scratch)
        """)


@responses.activate
def test_get_python_env_with_empty_response(runner: CliRunner):
    scratch_config = {
        "installed_packages": [],
    }
    response = responses.add(
        responses.GET,
        "http://localhost:8000/python_environment",
        json=scratch_config,
        status=200,
    )

    result = runner.invoke(main, ["controller", "get-python-env"])
    assert response.call_count == 1
    assert result.exit_code == 0
    assert result.output == dedent("""\
    Scratch Status: disabled
    No scratch packages found
    """)


def test_python_env_output_formatting():
    """Test for alternative python env output formats"""

    python_env_response = {
        "installed_packages": [
            {
                "name": "bar",
                "version": "0.0.1",
                "location": "/tmp/bar",
                "is_dirty": "false",
                "source": "pypi",
            },
            {
                "name": "foo",
                "version": "https://github.com/example/foo.git @18ec206e",
                "location": "/tmp/foo",
                "is_dirty": "true",
                "source": "scratch",
            },
        ],
        "scratch_enabled": "true",
    }
    python_env_response = PythonEnvironmentResponse(**python_env_response)

    compact = dedent("""\
        Scratch Status: enabled
        - bar @ (0.0.1)
        - foo @ (https://github.com/example/foo.git @18ec206e) (Dirty) (Scratch)
        """)

    _assert_matching_formatting(OutputFormat.COMPACT, python_env_response, compact)

    full = dedent("""\
        Scratch Status: enabled
        Installed Packages:
        - bar
        Version: 0.0.1
        Location: /tmp/bar
        Source: pypi
        Dirty: False
        - foo
        Version: https://github.com/example/foo.git @18ec206e
        Location: /tmp/foo
        Source: scratch
        Dirty: True
        """)

    _assert_matching_formatting(OutputFormat.FULL, python_env_response, full)

    empty_python_env = PythonEnvironmentResponse(
        installed_packages=[], scratch_enabled=False
    )

    full = dedent("""\
        Scratch Status: disabled
        No scratch packages found
        """)

    _assert_matching_formatting(OutputFormat.FULL, empty_python_env, full)


@pytest.mark.parametrize("output_flag", [True, False])
@pytest.mark.parametrize("update", [True, False])
@patch("blueapi.cli.cli.config.CONFIG_SCHEMA_LOCATION")
def test_config_schema(
    config_schema_location_mock: Mock,
    runner: CliRunner,
    output_flag: bool,
    update: bool,
    tmp_path: Path,
):
    args = ["config-schema"]

    tmp_path = tmp_path / "foo.json"

    if output_flag:
        args.append("-o")
        args.append(f"{tmp_path}")

    if update:
        args.append("-u")

    result = runner.invoke(
        main,
        args,
    )

    expected = generate_config_schema()
    if output_flag and (not update):
        with tmp_path.open("r") as stream:
            assert json.load(stream) == expected
    elif update:
        config_schema_location_mock.open.assert_called()
        with config_schema_location_mock.open() as stream:
            stream.write.assert_called()
    else:
        assert json.loads(result.output) == expected


@pytest.mark.parametrize("value,result", [({}, {}), ("{}", {}), (None, None)])
def test_task_parameter_type(value, result):
    t = ParametersType()
    assert t.convert(value, None, None) == result