From 7e0058a6117941f56c0de359703e673fe3572095 Mon Sep 17 00:00:00 2001
From: wyzula-jan <133381102+wyzula-jan@users.noreply.github.com>
Date: Thu, 29 Feb 2024 14:58:58 +0100
Subject: [PATCH] test(cli/generate_cli): test added

---
 tests/test_generate_cli_client.py | 79 +++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 tests/test_generate_cli_client.py

diff --git a/tests/test_generate_cli_client.py b/tests/test_generate_cli_client.py
new file mode 100644
index 00000000..59f067a9
--- /dev/null
+++ b/tests/test_generate_cli_client.py
@@ -0,0 +1,79 @@
+import pytest
+from bec_widgets.cli.generate_cli import ClientGenerator
+from textwrap import dedent
+import black
+
+
+# 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."""
+        pass
+
+    def remove_plot(self, plot_id: str):
+        """Remove a plot from the figure."""
+        pass
+
+
+def test_client_generator_with_black_formatting():
+    generator = ClientGenerator()
+    generator.generate_client([MockBECWaveform1D, MockBECFigure])
+
+    # 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 bec_widgets.cli.client_utils import rpc_call, RPCBase, BECFigureClientMixin
+        from typing import Literal, Optional, overload
+
+        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)
+    )
+
+    assert expected_output_formatted == generated_output_formatted