added tests

Author: twishabansal

packages/toolbox-langchain/tests/test_async_client.py

@@ -339,3 +339,22 @@ async def test_load_toolset_not_implemented(self, mock_client):
         assert "Synchronous methods not supported by async client." in str(
             excinfo.value
         )
+    
+    @patch("toolbox_langchain.async_client.ToolboxCoreClient")
+    async def test_init_with_client_headers(
+        self, mock_core_client_constructor, mock_session
+    ):
+        """Tests that client_headers are passed to the core client during initialization."""
+        headers = {"X-Test-Header": "value"}
+        AsyncToolboxClient(URL, session=mock_session, client_headers=headers)
+        mock_core_client_constructor.assert_called_once_with(
+            url=URL, session=mock_session, client_headers=headers
+        )
+
+    async def test_add_headers(self, mock_client):
+        """Tests that add_headers calls the core client's add_headers."""
+        headers = {"X-Another-Header": lambda: "dynamic_value"}
+        mock_client.add_headers(headers)
+        mock_client._AsyncToolboxClient__core_client.add_headers.assert_called_once_with(
+            headers
+        )

packages/toolbox-langchain/tests/test_client.py

@@ -420,3 +420,21 @@ async def test_aload_toolset_with_args(
             bound_params=bound_params,
             strict=True,
         )
+
+    @patch("toolbox_langchain.client.ToolboxCoreSyncClient")
+    def test_init_with_client_headers(self, mock_core_client_constructor):
+        """Tests that client_headers are passed to the core client during initialization."""
+        headers = {"X-Test-Header": "value"}
+        ToolboxClient(URL, client_headers=headers)
+        mock_core_client_constructor.assert_called_once_with(
+            url=URL, client_headers=headers
+        )
+
+    @patch("toolbox_langchain.client.ToolboxCoreSyncClient")
+    def test_add_headers(self, mock_core_client_constructor):
+        """Tests that add_headers calls the core client's add_headers."""
+        mock_core_instance = mock_core_client_constructor.return_value
+        client = ToolboxClient(URL)
+        headers = {"X-Another-Header": "dynamic_value"}
+        client.add_headers(headers)
+        mock_core_instance.add_headers.assert_called_once_with(headers)