0
0
mirror of https://github.com/bec-project/bec_widgets.git synced 2025-07-14 03:31:50 +02:00
Files
bec_widgets/tests/unit_tests/test_generate_cli_client.py

100 lines
2.7 KiB
Python

from textwrap import dedent
import black
import isort
from bec_widgets.cli.generate_cli import ClientGenerator
# 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()
rpc_classes = {
"connector_classes": [MockBECWaveform1D, MockBECFigure],
"top_level_classes": [MockBECFigure],
}
generator.generate_client(rpc_classes)
# 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
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 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.
"""
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.
"""
'''
)
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