feat: add comprehensive test coverage for tools module

README.md

@@ -528,18 +528,18 @@ When using the xAI SDK, you may encounter various error codes returned by the AP
 
 Below is a table of common gRPC status codes you might encounter when using the xAI SDK:
 
-| gRPC Status Code          | Meaning                                                                | xAI SDK/API Context                                                                                     |
-|---------------------------|------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------|
-| `UNKNOWN`                 | An unknown error occurred.                                             | An unexpected issue occurred on the server side, not specifically related to the request.               |
-| `INVALID_ARGUMENT`        | The client specified an invalid argument.                              | An invalid argument was provided to the model/endpoint, such as incorrect parameters or malformed input.|
-| `DEADLINE_EXCEEDED`       | The deadline for the request expired before the operation completed.   | Raised if the request exceeds the timeout specified by the client (default is 1620 seconds, configurable during client instantiation). |
-| `NOT_FOUND`               | A specified resource was not found.                                    | A requested model or resource does not exist.                                                           |
-| `PERMISSION_DENIED`       | The caller does not have permission to execute the specified operation.| The API key is disabled, blocked, or lacks sufficient permissions to access a specific model or feature. |
-| `UNAUTHENTICATED`         | The request does not have valid authentication credentials.            | The API key is missing, invalid, or expired.                                                            |
-| `RESOURCE_EXHAUSTED`      | A resource quota has been exceeded (e.g., rate limits).                | The user has exceeded their API usage quota or rate limits for requests.                                |
-| `INTERNAL`                | An internal error occurred.                                            | An internal server error occurred on the xAI API side.                                                  |
-| `UNAVAILABLE`             | The service is currently unavailable. This is often a transient error. | The model or endpoint invoked is temporarily down or there are connectivity issues. The SDK defaults to automatically retrying errors with this status code. |
-| `DATA_LOSS`               | Unrecoverable data loss or corruption occurred.                        | Occurs when a user provides an image via URL in API calls (e.g., in a chat conversation) and the server fails to fetch the image from that URL. |
+| gRPC Status Code      | Meaning                                                              | xAI SDK/API Context                                                                                        |
+|-----------------------|----------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------|
+| `UNKNOWN`             | An unknown error occurred.                                           | An unexpected issue occurred on the server side, not specifically related to the request.                  |
+| `INVALID_ARGUMENT`    | The client specified an invalid argument.                            | An invalid argument was provided to the model/endpoint, such as incorrect parameters or malformed input.   |
+| `DEADLINE_EXCEEDED`   | The deadline for the request expired before the operation completed. | Raised if the request exceeds the timeout specified by the client (default is 1620 seconds, configurable during client instantiation). |
+| `NOT_FOUND`           | A specified resource was not found.                                  | A requested model or resource does not exist.                                                              |
+| `PERMISSION_DENIED`   | The caller does not have permission to execute the specified operation. | The API key is disabled, blocked, or lacks sufficient permissions to access a specific model or feature.  |
+| `UNAUTHENTICATED`     | The request does not have valid authentication credentials.          | The API key is missing, invalid, or expired.                                                               |
+| `RESOURCE_EXHAUSTED`  | A resource quota has been exceeded (e.g., rate limits).              | The user has exceeded their API usage quota or rate limits for requests.                                   |
+| `INTERNAL`            | An internal error occurred.                                          | An internal server error occurred on the xAI API side.                                                     |
+| `UNAVAILABLE`         | The service is currently unavailable. This is often a transient error. | The model or endpoint invoked is temporarily down or there are connectivity issues. The SDK defaults to automatically retrying errors with this status code. |
+| `DATA_LOSS`           | Unrecoverable data loss or corruption occurred.                      | Occurs when a user provides an image via URL in API calls (e.g., in a chat conversation) and the server fails to fetch the image from that URL. |
 
 These error codes can help diagnose issues with API requests. When handling errors, ensure you check the specific status code to understand the nature of the problem and take appropriate action.
 

src/xai_sdk/tools.py

@@ -272,8 +272,6 @@ def mcp(
         server_url: The URL of the MCP server.
         server_label: Optional label of the MCP server. This will be used to prefix tool names if provided.
         server_description: Optional description of the MCP server.
-        server_label: The label of the MCP server. This will be used to prefix tool names if provided.
-        server_description: The description of the MCP server.
         allowed_tool_names: The names of the tools that the model is allowed to call. If empty, all tools are allowed.
         authorization: The authorization token for the MCP server.
         extra_headers: The extra headers for the MCP server.

tests/tools_test.py

@@ -0,0 +1,197 @@
+"""Tests for xai_sdk.tools module."""
+
+import datetime
+from unittest import mock
+
+import pytest
+
+from xai_sdk import tools
+from xai_sdk.proto import chat_pb2
+
+
+class TestWebSearch:
+    """Tests for web_search tool."""
+
+    def test_basic_web_search(self):
+        """Test creating a basic web search tool."""
+        tool = tools.web_search()
+        assert tool.HasField("web_search")
+        assert tool.web_search.excluded_domains == []
+        assert tool.web_search.allowed_domains == []
+
+    def test_web_search_with_excluded_domains(self):
+        """Test web search with excluded domains."""
+        excluded = ["spam.com", "fake.com"]
+        tool = tools.web_search(excluded_domains=excluded)
+        assert tool.web_search.excluded_domains == excluded
+
+    def test_web_search_with_allowed_domains(self):
+        """Test web search with allowed domains."""
+        allowed = ["github.com", "stackoverflow.com"]
+        tool = tools.web_search(allowed_domains=allowed)
+        assert tool.web_search.allowed_domains == allowed
+
+    def test_web_search_with_image_options(self):
+        """Test web search with image understanding and search enabled."""
+        tool = tools.web_search(
+            enable_image_understanding=True,
+            enable_image_search=True,
+        )
+        assert tool.web_search.enable_image_understanding
+        assert tool.web_search.enable_image_search
+
+    def test_web_search_with_location(self):
+        """Test web search with user location."""
+        tool = tools.web_search(
+            user_location_country="US",
+            user_location_city="San Francisco",
+            user_location_region="California",
+            user_location_timezone="America/Los_Angeles",
+        )
+        assert tool.web_search.user_location.country == "US"
+        assert tool.web_search.user_location.city == "San Francisco"
+        assert tool.web_search.user_location.region == "California"
+        assert tool.web_search.user_location.timezone == "America/Los_Angeles"
+
+
+class TestXSearch:
+    """Tests for x_search tool."""
+
+    def test_basic_x_search(self):
+        """Test creating a basic X search tool."""
+        tool = tools.x_search()
+        assert tool.HasField("x_search")
+
+    def test_x_search_with_date_range(self):
+        """Test X search with date range."""
+        from_date = datetime.datetime(2025, 1, 1)
+        to_date = datetime.datetime(2025, 1, 31)
+        tool = tools.x_search(from_date=from_date, to_date=to_date)
+        assert tool.x_search.from_date.seconds > 0
+        assert tool.x_search.to_date.seconds > 0
+
+    def test_x_search_with_handles(self):
+        """Test X search with allowed X handles."""
+        allowed = ["xai", "elonmusk"]
+        tool = tools.x_search(allowed_x_handles=allowed)
+        assert tool.x_search.allowed_x_handles == allowed
+
+    def test_x_search_with_excluded_handles(self):
+        """Test X search with excluded X handles."""
+        excluded = ["spam_bot", "fake_account"]
+        tool = tools.x_search(excluded_x_handles=excluded)
+        assert tool.x_search.excluded_x_handles == excluded
+
+    def test_x_search_with_media_options(self):
+        """Test X search with image and video understanding."""
+        tool = tools.x_search(
+            enable_image_understanding=True,
+            enable_video_understanding=True,
+        )
+        assert tool.x_search.enable_image_understanding
+        assert tool.x_search.enable_video_understanding
+
+
+class TestCodeExecution:
+    """Tests for code_execution tool."""
+
+    def test_code_execution(self):
+        """Test creating a code execution tool."""
+        tool = tools.code_execution()
+        assert tool.HasField("code_execution")
+
+
+class TestCollectionsSearch:
+    """Tests for collections_search tool."""
+
+    def test_basic_collections_search(self):
+        """Test creating a basic collections search tool."""
+        collection_ids = ["col_123", "col_456"]
+        tool = tools.collections_search(collection_ids)
+        assert tool.HasField("collections_search")
+        assert list(tool.collections_search.collection_ids) == collection_ids
+
+    def test_collections_search_with_limit(self):
+        """Test collections search with result limit."""
+        collection_ids = ["col_123"]
+        tool = tools.collections_search(collection_ids, limit=10)
+        assert tool.collections_search.limit == 10
+
+    def test_collections_search_with_instructions(self):
+        """Test collections search with custom instructions."""
+        collection_ids = ["col_123"]
+        instructions = "Find documents about machine learning"
+        tool = tools.collections_search(collection_ids, instructions=instructions)
+        assert tool.collections_search.instructions == instructions
+
+
+class TestMCP:
+    """Tests for mcp tool."""
+
+    def test_basic_mcp(self):
+        """Test creating a basic MCP tool."""
+        server_url = "http://localhost:3000"
+        tool = tools.mcp(server_url)
+        assert tool.HasField("mcp")
+        assert tool.mcp.server_url == server_url
+
+    def test_mcp_with_label_and_description(self):
+        """Test MCP tool with label and description."""
+        server_url = "http://localhost:3000"
+        label = "my_mcp_server"
+        description = "My custom MCP server"
+        tool = tools.mcp(server_url, server_label=label, server_description=description)
+        assert tool.mcp.server_label == label
+        assert tool.mcp.server_description == description
+
+    def test_mcp_with_allowed_tools(self):
+        """Test MCP tool with allowed tool names."""
+        server_url = "http://localhost:3000"
+        allowed_tools = ["tool1", "tool2", "tool3"]
+        tool = tools.mcp(server_url, allowed_tool_names=allowed_tools)
+        assert list(tool.mcp.allowed_tool_names) == allowed_tools
+
+    def test_mcp_with_authorization(self):
+        """Test MCP tool with authorization token."""
+        server_url = "http://localhost:3000"
+        auth = "Bearer token123"
+        tool = tools.mcp(server_url, authorization=auth)
+        assert tool.mcp.authorization == auth
+
+    def test_mcp_with_extra_headers(self):
+        """Test MCP tool with extra headers."""
+        server_url = "http://localhost:3000"
+        headers = {"X-Custom-Header": "value", "X-Another": "header"}
+        tool = tools.mcp(server_url, extra_headers=headers)
+        assert dict(tool.mcp.extra_headers) == headers
+
+
+class TestGetToolCallType:
+    """Tests for get_tool_call_type function."""
+
+    def test_get_web_search_tool_type(self):
+        """Test getting type of web search tool call."""
+        tool_call = chat_pb2.ToolCall(type=chat_pb2.TOOL_CALL_TYPE_WEB_SEARCH)
+        assert tools.get_tool_call_type(tool_call) == "web_search_tool"
+
+    def test_get_x_search_tool_type(self):
+        """Test getting type of X search tool call."""
+        tool_call = chat_pb2.ToolCall(type=chat_pb2.TOOL_CALL_TYPE_X_SEARCH)
+        assert tools.get_tool_call_type(tool_call) == "x_search_tool"
+
+    def test_get_code_execution_tool_type(self):
+        """Test getting type of code execution tool call."""
+        tool_call = chat_pb2.ToolCall(type=chat_pb2.TOOL_CALL_TYPE_CODE_EXECUTION)
+        assert tools.get_tool_call_type(tool_call) == "code_execution_tool"
+
+    def test_get_collections_search_tool_type(self):
+        """Test getting type of collections search tool call."""
+        tool_call = chat_pb2.ToolCall(
+            type=chat_pb2.TOOL_CALL_TYPE_COLLECTIONS_SEARCH
+        )
+        assert tools.get_tool_call_type(tool_call) == "collections_search_tool"
+
+    def test_get_mcp_tool_type(self):
+        """Test getting type of MCP tool call."""
+        tool_call = chat_pb2.ToolCall(type=chat_pb2.TOOL_CALL_TYPE_MCP)
+        assert tools.get_tool_call_type(tool_call) == "mcp_tool"