bec_widgets/tests/unit_tests/test_generate_cli_client.py

from textwrap import dedent

import black
import isort

from bec_widgets.cli.generate_cli import ClientGenerator
from bec_widgets.utils.plugin_utils import BECClassContainer, BECClassInfo

# pylint: disable=missing-function-docstring


# Mock classes to test the generator
class MockBECWaveform1D:
    USER_ACCESS = ["set_frequency", "set_amplitude"]

    def set_frequency(self, frequency: float) -> list:
        """Set the frequency of the waveform."""
        return [frequency]

    def set_amplitude(self, amplitude: float) -> tuple[float, float]:
        """Set the amplitude of the waveform."""
        return amplitude, amplitude


class MockBECFigure:
    USER_ACCESS = ["add_plot", "remove_plot"]

    def add_plot(self, plot_id: str):
        """Add a plot to the figure."""

    def remove_plot(self, plot_id: str):
        """Remove a plot from the figure."""


def test_client_generator_with_black_formatting():
    generator = ClientGenerator()
    container = BECClassContainer()
    container.add_class(
        BECClassInfo(
            name="MockBECWaveform1D",
            module="test_module",
            file="test_file",
            obj=MockBECWaveform1D,
            is_connector=True,
            is_widget=True,
            is_plugin=False,
        )
    )
    container.add_class(
        BECClassInfo(
            name="MockBECFigure",
            module="test_module",
            file="test_file",
            obj=MockBECFigure,
            is_connector=True,
            is_widget=True,
            is_plugin=True,
        )
    )

    generator.generate_client(container)

    # Format the expected output with black to ensure it matches the generator output
    expected_output = dedent(
        '''\
        # This file was automatically generated by generate_cli.py

        from __future__ import annotations

        import enum
        from typing import Literal, Optional, overload

        from bec_widgets.cli.client_utils import BECGuiClientMixin, RPCBase, rpc_call

        # pylint: skip-file


        class Widgets(str, enum.Enum):
            """
            Enum for the available widgets.
            """

            MockBECFigure = "MockBECFigure"


        class MockBECFigure(RPCBase):
            @rpc_call
            def add_plot(self, plot_id: str):
                """
                Add a plot to the figure.
                """

            @rpc_call
            def remove_plot(self, plot_id: str):
                """
                Remove a plot from the figure.
                """


        class MockBECWaveform1D(RPCBase):
            @rpc_call
            def set_frequency(self, frequency: float) -> list:
                """
                Set the frequency of the waveform.
                """

            @rpc_call
            def set_amplitude(self, amplitude: float) -> tuple[float, float]:
                """
                Set the amplitude of the waveform.
                """
    '''
    )

    expected_output_formatted = black.format_str(
        expected_output, mode=black.FileMode(line_length=100)
    ).lstrip()

    generated_output_formatted = black.format_str(
        generator.header + "\n" + generator.content, mode=black.FileMode(line_length=100)
    )

    generated_output_formatted = isort.code(generated_output_formatted)

    assert expected_output_formatted == generated_output_formatted