feat: updated tests searchscraper

PeriniM · PeriniM · commit a771564838b6 · 2025-02-23T00:08:44.000+01:00
diff --git a/tests/unit_tests/mocks.py b/tests/unit_tests/mocks.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Type
 
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel, Field
@@ -115,6 +115,7 @@ class MockSearchScraperTool(BaseTool):
     args_schema: type[BaseModel] = MockSearchScraperInput
     client: Optional[MockClient] = None
     api_key: str
+    llm_output_schema: Optional[Type[BaseModel]] = None
 
     def _run(self, **kwargs: Any) -> Dict:
         return {
diff --git a/tests/unit_tests/test_tools.py b/tests/unit_tests/test_tools.py
@@ -35,19 +35,22 @@ def tool_invoke_params_example(self) -> dict:
             "website_url": "https://example.com",
         }
 
+
+class TestSmartScraperToolCustom:
     def test_invoke_with_html(self):
         """Test invoking the tool with HTML content."""
-        tool = self.tool_constructor(**self.tool_constructor_params)
-        result = tool.invoke(
-            {
-                "user_prompt": "Extract the main heading",
-                "website_url": "https://example.com",
-                "website_html": "<html><body><h1>Test</h1></body></html>",
-            }
-        )
-        assert isinstance(result, dict)
-        assert "main_heading" in result
-        assert result["main_heading"] == "Test"
+        with patch("langchain_scrapegraph.tools.smartscraper.Client", MockClient):
+            tool = MockSmartScraperTool(api_key="sgai-test-api-key")
+            result = tool.invoke(
+                {
+                    "user_prompt": "Extract the main heading",
+                    "website_url": "https://example.com",
+                    "website_html": "<html><body><h1>Test</h1></body></html>",
+                }
+            )
+            assert isinstance(result, dict)
+            assert "main_heading" in result
+            assert result["main_heading"] == "Test"
 
 
 class TestSearchScraperToolUnit(ToolsUnitTests):
@@ -66,6 +69,8 @@ def tool_invoke_params_example(self) -> dict:
             "user_prompt": "What are the key features of Product X?",
         }
 
+
+class TestSearchScraperToolCustom:
     def test_invoke_with_schema(self):
         """Test invoking the tool with a schema."""
         from typing import List
@@ -77,14 +82,17 @@ class TestSchema(BaseModel):
             features: List[dict] = Field(description="List of features")
             reference_urls: List[str] = Field(description="Reference URLs")
 
-        tool = self.tool_constructor(**self.tool_constructor_params)
-        tool.llm_output_schema = TestSchema
-        result = tool.invoke(self.tool_invoke_params_example)
-        assert isinstance(result, dict)
-        assert "product" in result
-        assert "features" in result
-        assert "reference_urls" in result
-        assert isinstance(result["reference_urls"], list)
+        with patch("langchain_scrapegraph.tools.searchscraper.Client", MockClient):
+            tool = MockSearchScraperTool(api_key="sgai-test-api-key")
+            tool.llm_output_schema = TestSchema
+            result = tool.invoke(
+                {"user_prompt": "What are the key features of Product X?"}
+            )
+            assert isinstance(result, dict)
+            assert "product" in result
+            assert "features" in result
+            assert "reference_urls" in result
+            assert isinstance(result["reference_urls"], list)
 
 
 class TestGetCreditsToolUnit(ToolsUnitTests):
