From 9539edddaa24a119c7e12c1279c25b861bf3ca23 Mon Sep 17 00:00:00 2001
From: Luca Beurer-Kellner <lucabeurerkellner@gmail.com>
Date: Tue, 6 May 2025 17:44:53 +0200
Subject: [PATCH] extra metadata tests

---
 run.sh                                               |  3 +++
 tests/integration/mcp/test_mcp_stdio.py              |  5 +++++
 tests/integration/resources/mcp/stdio/client/main.py | 10 ++++++++++
 3 files changed, 18 insertions(+)

diff --git a/run.sh b/run.sh
index 0b95bd8..8ae0930 100755
--- a/run.sh
+++ b/run.sh
@@ -136,6 +136,9 @@ integration_tests() {
     sleep 2
   done
 
+  # Delete existing dist files
+  rm -rf dist
+
   # Generate latest whl file for the invariant-gateway package. 
   # This is required to run the integration tests.
   pip install build
diff --git a/tests/integration/mcp/test_mcp_stdio.py b/tests/integration/mcp/test_mcp_stdio.py
index 8946cad..9313015 100644
--- a/tests/integration/mcp/test_mcp_stdio.py
+++ b/tests/integration/mcp/test_mcp_stdio.py
@@ -28,6 +28,7 @@ async def test_mcp_stdio_with_gateway(
         push_to_explorer=push_to_explorer,
         tool_name="get_last_message_from_user",
         tool_args={"username": "Alice"},
+        metadata_keys={"my-custom-key": "value1", "my-custom-key-2": "value2"},
     )
 
     assert result.isError is False
@@ -58,6 +59,10 @@ async def test_mcp_stdio_with_gateway(
             and metadata["mcp_client"] == "mcp"
             and metadata["mcp_server"] == "messenger_server"
         )
+        # ensure custom keys are present
+        assert metadata["my-custom-key"] == "value1"
+        assert metadata["my-custom-key-2"] == "value2"
+        
         assert trace["messages"][0]["role"] == "assistant"
         assert trace["messages"][0]["tool_calls"][0]["function"] == {
             "name": "get_last_message_from_user",
diff --git a/tests/integration/resources/mcp/stdio/client/main.py b/tests/integration/resources/mcp/stdio/client/main.py
index 4238949..56a7f65 100644
--- a/tests/integration/resources/mcp/stdio/client/main.py
+++ b/tests/integration/resources/mcp/stdio/client/main.py
@@ -23,6 +23,7 @@ class MCPClient:
         project_name: str,
         server_script_path: str,
         push_to_explorer: bool,
+        metadata_keys: Optional[dict[str, str]] = None,
     ):
         """
         Connect to an MCP server.
@@ -42,6 +43,12 @@ class MCPClient:
             "--project-name",
             project_name,
         ]
+        # add metadata cli args
+        if metadata_keys is not None:
+            for key, value in metadata_keys.items():
+                args.append("--metadata-" + key + "=" + value)
+
+
         if push_to_explorer:
             args.append("--push-explorer")
         args.extend(
@@ -54,6 +61,7 @@ class MCPClient:
                 os.path.basename(server_script_path),
             ]
         )
+
         server_params = StdioServerParameters(
             command="uvx",
             args=args,
@@ -103,6 +111,7 @@ async def run(
     push_to_explorer: bool,
     tool_name: str,
     tool_args: dict[str, Any],
+    metadata_keys: Optional[dict[str, str]] = None,
 ) -> types.CallToolResult:
     """
     Main function to setup the MCP client and server.
@@ -125,6 +134,7 @@ async def run(
             project_name,
             server_script_path,
             push_to_explorer,
+            metadata_keys=metadata_keys
         )
         return await client.call_tool(tool_name, tool_args)
     finally: