From c5b4be695019032e246d01e77d53a43989a4a7ec Mon Sep 17 00:00:00 2001
From: Yoel Bassin <joel@Mac-2.lan>
Date: Sat, 26 Apr 2025 16:40:39 +0300
Subject: [PATCH] main - feat: Implement flowgraph block removal

---
 src/gnuradio_mcp/middlewares/flowgraph.py | 11 +++--
 tests/test_flowgraph.py                   | 54 ++++++++++++++++++++---
 2 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/src/gnuradio_mcp/middlewares/flowgraph.py b/src/gnuradio_mcp/middlewares/flowgraph.py
index 42ab8f6..1694240 100644
--- a/src/gnuradio_mcp/middlewares/flowgraph.py
+++ b/src/gnuradio_mcp/middlewares/flowgraph.py
@@ -15,6 +15,11 @@ class FlowGraphMiddleware:
             for block in self._flowgraph.blocks
         ]
     
-    def add_block(self, key: str, name) -> None:
-        block = self._flowgraph.new_block(key)
-        block.params["id"].set_value(name)
+    def add_block(self, block_type: str, block_name: str) -> None:
+        block = self._flowgraph.new_block(block_type)
+        block.params["id"].set_value(block_name)
+
+
+    def remove_block(self, block_name: str) -> None:
+        block = self._flowgraph.get_block(block_name)
+        self._flowgraph.remove_element(block)
diff --git a/tests/test_flowgraph.py b/tests/test_flowgraph.py
index d16ad62..79b3235 100644
--- a/tests/test_flowgraph.py
+++ b/tests/test_flowgraph.py
@@ -5,6 +5,9 @@ from gnuradio import gr
 
 from gnuradio_mcp.models import BlockModel
 
+BLOCK_KEYS_TO_TEST = [1, 10]
+
+
 @pytest.fixture(scope="module")
 def platform() -> Platform:
     platform = Platform(
@@ -15,11 +18,13 @@ def platform() -> Platform:
     platform.build_library()
     return platform
 
+
 @pytest.fixture
 def flowgraph_middleware(platform):
     flowgraph = platform.make_flow_graph("")
     return FlowGraphMiddleware(flowgraph)
 
+
 @pytest.fixture
 def initial_blocks(flowgraph_middleware):
     return [
@@ -27,16 +32,55 @@ def initial_blocks(flowgraph_middleware):
         for block in flowgraph_middleware._flowgraph.blocks
     ]
 
+
+def get_block_keys(platform):
+    return list(platform.blocks.keys())
+
+
+def add_block(flowgraph_middleware, block_key, block_name):
+    flowgraph_middleware.add_block(block_key, block_name)
+
+
+def remove_block(flowgraph_middleware, block_name):
+    flowgraph_middleware._flowgraph.remove_element(
+        flowgraph_middleware._flowgraph.get_block(block_name)
+    )
+
+
+def get_current_blocks(flowgraph_middleware):
+    return [
+        BlockModel(key=block.key, label=block.label)
+        for block in flowgraph_middleware._flowgraph.blocks
+    ]
+
+
 def test_blocks_match_initial_state(flowgraph_middleware, initial_blocks):
     assert flowgraph_middleware.blocks == initial_blocks
 
-@pytest.mark.parametrize("block_index", [1, 2, 3])
-def test_add_block_preserves_and_adds(flowgraph_middleware, platform, initial_blocks, block_index):
-    block_keys = list(platform.blocks.keys())
+
+@pytest.mark.parametrize("block_index", BLOCK_KEYS_TO_TEST)
+def test_add_block_preserves_and_adds(
+    flowgraph_middleware, platform, initial_blocks, block_index
+):
+    block_keys = get_block_keys(platform)
     assert block_keys, "No blocks available in platform library."
     block_key = block_keys[block_index]
-    block_name = "test_block"
-    flowgraph_middleware.add_block(block_key, block_name)
+    block_name = f"test_block_{block_index}"
+    add_block(flowgraph_middleware, block_key, block_name)
     blocks = flowgraph_middleware.blocks
     assert all(b in blocks for b in initial_blocks)
     assert any(b.key == block_key for b in blocks)
+
+
+@pytest.mark.parametrize("block_index", BLOCK_KEYS_TO_TEST)
+def test_remove_block_restores_initial(
+    flowgraph_middleware, platform, initial_blocks, block_index
+):
+    block_keys = get_block_keys(platform)
+    block_key = block_keys[block_index]
+    block_name = f"block_to_remove_{block_index}"
+    add_block(flowgraph_middleware, block_key, block_name)
+    assert any(b.key == block_key for b in flowgraph_middleware.blocks)
+    remove_block(flowgraph_middleware, block_name)
+    current_blocks = get_current_blocks(flowgraph_middleware)
+    assert current_blocks == initial_blocks