Adding Snowflake search tool

tests/tools/snowflake_search_tool_test.py

@@ -0,0 +1,103 @@
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from crewai_tools import SnowflakeConfig, SnowflakeSearchTool
+
+
+# Unit Test Fixtures
+@pytest.fixture
+def mock_snowflake_connection():
+    mock_conn = MagicMock()
+    mock_cursor = MagicMock()
+    mock_cursor.description = [("col1",), ("col2",)]
+    mock_cursor.fetchall.return_value = [(1, "value1"), (2, "value2")]
+    mock_cursor.execute.return_value = None
+    mock_conn.cursor.return_value = mock_cursor
+    return mock_conn
+
+
+@pytest.fixture
+def mock_config():
+    return SnowflakeConfig(
+        account="test_account",
+        user="test_user",
+        password="test_password",
+        warehouse="test_warehouse",
+        database="test_db",
+        snowflake_schema="test_schema",
+    )
+
+
+@pytest.fixture
+def snowflake_tool(mock_config):
+    with patch("snowflake.connector.connect") as mock_connect:
+        tool = SnowflakeSearchTool(config=mock_config)
+        yield tool
+
+
+# Unit Tests
+@pytest.mark.asyncio
+async def test_successful_query_execution(snowflake_tool, mock_snowflake_connection):
+    with patch.object(snowflake_tool, "_create_connection") as mock_create_conn:
+        mock_create_conn.return_value = mock_snowflake_connection
+
+        results = await snowflake_tool._run(
+            query="SELECT * FROM test_table", timeout=300
+        )
+
+        assert len(results) == 2
+        assert results[0]["col1"] == 1
+        assert results[0]["col2"] == "value1"
+        mock_snowflake_connection.cursor.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_connection_pooling(snowflake_tool, mock_snowflake_connection):
+    with patch.object(snowflake_tool, "_create_connection") as mock_create_conn:
+        mock_create_conn.return_value = mock_snowflake_connection
+
+        # Execute multiple queries
+        await asyncio.gather(
+            snowflake_tool._run("SELECT 1"),
+            snowflake_tool._run("SELECT 2"),
+            snowflake_tool._run("SELECT 3"),
+        )
+
+        # Should reuse connections from pool
+        assert mock_create_conn.call_count <= snowflake_tool.pool_size
+
+
+@pytest.mark.asyncio
+async def test_cleanup_on_deletion(snowflake_tool, mock_snowflake_connection):
+    with patch.object(snowflake_tool, "_create_connection") as mock_create_conn:
+        mock_create_conn.return_value = mock_snowflake_connection
+
+        # Add connection to pool
+        await snowflake_tool._get_connection()
+
+        # Return connection to pool
+        async with snowflake_tool._pool_lock:
+            snowflake_tool._connection_pool.append(mock_snowflake_connection)
+
+        # Trigger cleanup
+        snowflake_tool.__del__()
+
+        mock_snowflake_connection.close.assert_called_once()
+
+
+def test_config_validation():
+    # Test missing required fields
+    with pytest.raises(ValueError):
+        SnowflakeConfig()
+
+    # Test invalid account format
+    with pytest.raises(ValueError):
+        SnowflakeConfig(
+            account="invalid//account", user="test_user", password="test_pass"
+        )
+
+    # Test missing authentication
+    with pytest.raises(ValueError):
+        SnowflakeConfig(account="test_account", user="test_user")

tests/tools/test_code_interpreter_tool.py

@@ -7,7 +7,9 @@ from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import (
 
 
 class TestCodeInterpreterTool(unittest.TestCase):
-    @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
+    @patch(
+        "crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env"
+    )
     def test_run_code_in_docker(self, docker_mock):
         tool = CodeInterpreterTool()
         code = "print('Hello, World!')"
@@ -15,14 +17,14 @@ class TestCodeInterpreterTool(unittest.TestCase):
         expected_output = "Hello, World!\n"
 
         docker_mock().containers.run().exec_run().exit_code = 0
-        docker_mock().containers.run().exec_run().output = (
-            expected_output.encode()
-        )
+        docker_mock().containers.run().exec_run().output = expected_output.encode()
         result = tool.run_code_in_docker(code, libraries_used)
 
         self.assertEqual(result, expected_output)
 
-    @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
+    @patch(
+        "crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env"
+    )
     def test_run_code_in_docker_with_error(self, docker_mock):
         tool = CodeInterpreterTool()
         code = "print(1/0)"
@@ -37,7 +39,9 @@ class TestCodeInterpreterTool(unittest.TestCase):
 
         self.assertEqual(result, expected_output)
 
-    @patch("crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env")
+    @patch(
+        "crewai_tools.tools.code_interpreter_tool.code_interpreter_tool.docker_from_env"
+    )
     def test_run_code_in_docker_with_script(self, docker_mock):
         tool = CodeInterpreterTool()
         code = """print("This is line 1")