From e1282217c80e383c182e4e123e8512c377cc977c Mon Sep 17 00:00:00 2001 From: Gabriele Venturi Date: Fri, 31 Jan 2025 16:49:06 +0100 Subject: [PATCH 1/6] test: add more tests in the agent --- tests/unit_tests/agent/test_agent.py | 109 ++++++++++++++++++++++++++- 1 file changed, 107 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/agent/test_agent.py b/tests/unit_tests/agent/test_agent.py index f92a2316a..87edf1889 100644 --- a/tests/unit_tests/agent/test_agent.py +++ b/tests/unit_tests/agent/test_agent.py @@ -1,6 +1,6 @@ import os from typing import Optional -from unittest.mock import MagicMock, Mock, mock_open, patch +from unittest.mock import ANY, MagicMock, Mock, mock_open, patch import pandas as pd import pytest @@ -8,9 +8,10 @@ from pandasai import DatasetLoader, VirtualDataFrame from pandasai.agent.base import Agent from pandasai.config import Config, ConfigManager +from pandasai.core.response.error import ErrorResponse from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema from pandasai.dataframe.base import DataFrame -from pandasai.exceptions import CodeExecutionError +from pandasai.exceptions import CodeExecutionError, InvalidLLMOutputType from pandasai.llm.fake import FakeLLM @@ -466,3 +467,107 @@ def test_execute_sql_query_error_no_dataframe(self, agent): with pytest.raises(ValueError, match="No DataFrames available"): agent._execute_sql_query(query) + + def test_process_query(self, agent, config): + """Test the _process_query method with successful execution""" + query = "What is the average age?" + output_type = "number" + + # Mock the necessary methods + agent.generate_code = Mock(return_value="result = df['age'].mean()") + agent.execute_with_retries = Mock(return_value=30.5) + agent._state.config.enable_cache = True + agent._state.cache = Mock() + + # Execute the query + result = agent._process_query(query, output_type) + + # Verify the result + assert result == 30.5 + + # Verify method calls + agent.generate_code.assert_called_once() + agent.execute_with_retries.assert_called_once_with("result = df['age'].mean()") + agent._state.cache.set.assert_called_once() + + def test_process_query_execution_error(self, agent, config): + """Test the _process_query method with execution error""" + query = "What is the invalid operation?" + + # Mock methods to simulate error + agent.generate_code = Mock(return_value="invalid_code") + agent.execute_with_retries = Mock( + side_effect=CodeExecutionError("Execution failed") + ) + agent._handle_exception = Mock(return_value="Error handled") + + # Execute the query + result = agent._process_query(query) + + # Verify error handling + assert result == "Error handled" + agent._handle_exception.assert_called_once_with("invalid_code") + + def test_regenerate_code_after_invalid_llm_output_error(self, agent): + """Test code regeneration with InvalidLLMOutputType error""" + from pandasai.exceptions import InvalidLLMOutputType + + code = "test code" + error = InvalidLLMOutputType("Invalid output type") + + with patch( + "pandasai.agent.base.get_correct_output_type_error_prompt" + ) as mock_prompt: + mock_prompt.return_value = "corrected prompt" + agent._code_generator.generate_code = MagicMock(return_value="new code") + + result = agent._regenerate_code_after_error(code, error) + + mock_prompt.assert_called_once_with(agent._state, code, ANY) + agent._code_generator.generate_code.assert_called_once_with( + "corrected prompt" + ) + assert result == "new code" + + def test_regenerate_code_after_other_error(self, agent): + """Test code regeneration with non-InvalidLLMOutputType error""" + code = "test code" + error = ValueError("Some other error") + + with patch( + "pandasai.agent.base.get_correct_error_prompt_for_sql" + ) as mock_prompt: + mock_prompt.return_value = "sql error prompt" + agent._code_generator.generate_code = MagicMock(return_value="new code") + + result = agent._regenerate_code_after_error(code, error) + + mock_prompt.assert_called_once_with(agent._state, code, ANY) + agent._code_generator.generate_code.assert_called_once_with( + "sql error prompt" + ) + assert result == "new code" + + def test_handle_exception(self, agent): + """Test that _handle_exception properly formats and logs exceptions""" + test_code = "print(1/0)" # Code that will raise a ZeroDivisionError + + # Mock the logger to verify it's called + mock_logger = MagicMock() + agent._state.logger = mock_logger + + # Create an actual exception to handle + try: + exec(test_code) + except: + # Call the method + result = agent._handle_exception(test_code) + + # Verify the result is an ErrorResponse + assert isinstance(result, ErrorResponse) + assert result.last_code_executed == test_code + assert "ZeroDivisionError" in result.error + + # Verify the error was logged + mock_logger.log.assert_called_once() + assert "Processing failed with error" in mock_logger.log.call_args[0][0] From c9c72e732d58d3afc568264015eeb09786aa07c2 Mon Sep 17 00:00:00 2001 From: Gabriele Venturi Date: Fri, 31 Jan 2025 17:00:39 +0100 Subject: [PATCH 2/6] test: add more tests for the logger --- tests/unit_tests/helpers/test_logger.py | 56 +++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 tests/unit_tests/helpers/test_logger.py diff --git a/tests/unit_tests/helpers/test_logger.py b/tests/unit_tests/helpers/test_logger.py new file mode 100644 index 000000000..a47a9a907 --- /dev/null +++ b/tests/unit_tests/helpers/test_logger.py @@ -0,0 +1,56 @@ +import logging +from pandasai.helpers.logger import Logger + + +def test_verbose_setter(): + # Initialize logger with verbose=False + logger = Logger(verbose=False) + assert logger._verbose is False + assert not any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + + # Set verbose to True + logger.verbose = True + assert logger._verbose is True + assert any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert len(logger._logger.handlers) == 1 + + # Set verbose to False + logger.verbose = False + assert logger._verbose is False + assert not any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert len(logger._logger.handlers) == 0 + + # Set verbose to True again to ensure multiple toggles work + logger.verbose = True + assert logger._verbose is True + assert any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert len(logger._logger.handlers) == 1 + +def test_save_logs_property(): + # Initialize logger with save_logs=False + logger = Logger(save_logs=False, verbose=False) + assert logger.save_logs is False + + # Enable save_logs + logger.save_logs = True + assert logger.save_logs is True + assert any(isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers) + + # Disable save_logs + logger.save_logs = False + assert logger.save_logs is False + assert not any(isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers) + +def test_save_logs_property(): + # When logger is initialized with save_logs=True (default), it should have handlers + logger = Logger(save_logs=True) + assert logger.save_logs is True + + # When logger is initialized with save_logs=False, it should still have handlers if verbose=True + logger = Logger(save_logs=False, verbose=True) + assert logger.save_logs is True + + # When both save_logs and verbose are False, there should be no handlers + logger = Logger(save_logs=False, verbose=False) + logger._logger.handlers = [] # Reset handlers to match the property's expected behavior + assert logger.save_logs is False From 94a46419b892dce1e002b43c450afe7ba0253bd3 Mon Sep 17 00:00:00 2001 From: Gabriele Venturi Date: Fri, 31 Jan 2025 17:04:37 +0100 Subject: [PATCH 3/6] test: add more tests for the memory --- tests/test_memory.py | 80 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 tests/test_memory.py diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 000000000..15d3ebec9 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,80 @@ +import pytest +from pandasai.helpers.memory import Memory + + +def test_to_json_empty_memory(): + memory = Memory() + assert memory.to_json() == [] + + +def test_to_json_with_messages(): + memory = Memory() + + # Add test messages + memory.add("Hello", is_user=True) + memory.add("Hi there!", is_user=False) + memory.add("How are you?", is_user=True) + + expected_json = [ + {"role": "user", "message": "Hello"}, + {"role": "assistant", "message": "Hi there!"}, + {"role": "user", "message": "How are you?"} + ] + + assert memory.to_json() == expected_json + + +def test_to_json_message_order(): + memory = Memory() + + # Add messages in specific order + messages = [ + ("Message 1", True), + ("Message 2", False), + ("Message 3", True) + ] + + for msg, is_user in messages: + memory.add(msg, is_user=is_user) + + result = memory.to_json() + + # Verify order is preserved + assert len(result) == 3 + assert result[0]["message"] == "Message 1" + assert result[1]["message"] == "Message 2" + assert result[2]["message"] == "Message 3" + + +def test_to_openai_messages_empty(): + memory = Memory() + assert memory.to_openai_messages() == [] + + +def test_to_openai_messages_with_agent_description(): + memory = Memory(agent_description="I am a helpful assistant") + memory.add("Hello", is_user=True) + memory.add("Hi there!", is_user=False) + + expected_messages = [ + {"role": "system", "content": "I am a helpful assistant"}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"} + ] + + assert memory.to_openai_messages() == expected_messages + + +def test_to_openai_messages_without_agent_description(): + memory = Memory() + memory.add("Hello", is_user=True) + memory.add("Hi there!", is_user=False) + memory.add("How are you?", is_user=True) + + expected_messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + {"role": "user", "content": "How are you?"} + ] + + assert memory.to_openai_messages() == expected_messages From b0f98a9430bdc9ac75496b792d02f0b5173fc151 Mon Sep 17 00:00:00 2001 From: Gabriele Venturi Date: Fri, 31 Jan 2025 17:10:21 +0100 Subject: [PATCH 4/6] test: add more tests for the session --- tests/unit_tests/helpers/test_session.py | 76 +++++++++++++++++++++++- 1 file changed, 75 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/helpers/test_session.py b/tests/unit_tests/helpers/test_session.py index 2baeae5c0..264766ab0 100644 --- a/tests/unit_tests/helpers/test_session.py +++ b/tests/unit_tests/helpers/test_session.py @@ -2,9 +2,10 @@ from unittest.mock import patch import pytest +import requests from pandasai.constants import DEFAULT_API_URL -from pandasai.exceptions import PandaAIApiKeyError +from pandasai.exceptions import PandaAIApiKeyError, PandaAIApiCallError from pandasai.helpers.session import Session, get_pandaai_session @@ -107,3 +108,76 @@ def test_get_pandaai_session_with_env_api_url(): """Test that get_pandaai_session uses URL from environment""" session = get_pandaai_session() assert session._endpoint_url == "https://env.api.url" + + +@patch("requests.request") +def test_make_request_success(mock_request): + """Test successful API request""" + # Mock successful response + mock_response = mock_request.return_value + mock_response.status_code = 200 + mock_response.json.return_value = {"data": "test_data"} + + session = Session(api_key="test-key") + result = session.make_request("GET", "/test") + + # Verify request was made correctly + mock_request.assert_called_once_with( + "GET", + DEFAULT_API_URL + "/api/test", + headers={ + "x-authorization": "Bearer test-key", + "Content-Type": "application/json", + }, + params=None, + data=None, + json=None, + timeout=300, + ) + assert result == {"data": "test_data"} + + +@patch("requests.request") +def test_make_request_error_response(mock_request): + """Test API request with error response""" + # Mock error response + mock_response = mock_request.return_value + mock_response.status_code = 400 + mock_response.json.return_value = {"message": "Bad request"} + + session = Session(api_key="test-key") + with pytest.raises(PandaAIApiCallError) as exc_info: + session.make_request("POST", "/test") + + assert str(exc_info.value) == "Bad request" + + +@patch("requests.request") +def test_make_request_network_error(mock_request): + """Test API request with network error""" + # Mock network error + mock_request.side_effect = requests.exceptions.RequestException("Network error") + + session = Session(api_key="test-key") + with pytest.raises(PandaAIApiCallError) as exc_info: + session.make_request("GET", "/test") + + assert "Request failed: Network error" in str(exc_info.value) + + +@patch("requests.request") +def test_make_request_custom_headers(mock_request): + """Test API request with custom headers""" + # Mock successful response + mock_response = mock_request.return_value + mock_response.status_code = 200 + mock_response.json.return_value = {"data": "test_data"} + + custom_headers = {"Custom-Header": "test-value"} + session = Session(api_key="test-key") + session.make_request("GET", "/test", headers=custom_headers) + + # Verify custom headers were used + called_headers = mock_request.call_args[1]["headers"] + assert called_headers["Custom-Header"] == "test-value" + assert "x-authorization" not in called_headers From a0289ee2e22cec9dd9030b2c786a634cb6f720a1 Mon Sep 17 00:00:00 2001 From: Gabriele Venturi Date: Fri, 31 Jan 2025 17:16:06 +0100 Subject: [PATCH 5/6] test: add more tests for the pull --- tests/unit_tests/dataframe/test_pull.py | 119 ++++++++++++++++++++++++ 1 file changed, 119 insertions(+) create mode 100644 tests/unit_tests/dataframe/test_pull.py diff --git a/tests/unit_tests/dataframe/test_pull.py b/tests/unit_tests/dataframe/test_pull.py new file mode 100644 index 000000000..d6da922fa --- /dev/null +++ b/tests/unit_tests/dataframe/test_pull.py @@ -0,0 +1,119 @@ +import os +import pytest +from unittest.mock import patch, Mock, mock_open +from io import BytesIO +from zipfile import ZipFile + +import pandas as pd +from pandasai.dataframe.base import DataFrame +from pandasai.exceptions import PandaAIApiKeyError, DatasetNotFound +from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema, Column, Source + + +@pytest.fixture +def mock_env(monkeypatch): + monkeypatch.setenv("PANDABI_API_KEY", "test_api_key") + + +@pytest.fixture +def sample_df(): + return pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) + + +@pytest.fixture +def mock_zip_content(): + zip_buffer = BytesIO() + with ZipFile(zip_buffer, "w") as zip_file: + zip_file.writestr("test.csv", "col1,col2\n1,a\n2,b\n3,c") + return zip_buffer.getvalue() + + +@pytest.fixture +def mock_schema(): + return SemanticLayerSchema( + name="test_schema", + source=Source(type="parquet", path="data.parquet", table="test_table"), + columns=[ + Column(name="col1", type="integer"), + Column(name="col2", type="string"), + ], + ) + + +def test_pull_success(mock_env, sample_df, mock_zip_content, mock_schema, tmp_path): + with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, \ + patch("pandasai.dataframe.base.find_project_root") as mock_root, \ + patch("pandasai.DatasetLoader.create_loader_from_path") as mock_loader, \ + patch("builtins.open", mock_open()) as mock_file: + + # Setup mocks + mock_response = Mock() + mock_response.status_code = 200 + mock_response.content = mock_zip_content + mock_session.return_value.get.return_value = mock_response + mock_root.return_value = str(tmp_path) + + mock_loader_instance = Mock() + mock_loader_instance.load.return_value = DataFrame(sample_df, schema=mock_schema) + mock_loader.return_value = mock_loader_instance + + # Create DataFrame instance and call pull + df = DataFrame(sample_df, path="test/path", schema=mock_schema) + df.pull() + + # Verify API call + mock_session.return_value.get.assert_called_once_with( + "/datasets/pull", + headers={"accept": "application/json", "x-authorization": "Bearer test_api_key"}, + params={"path": "test/path"} + ) + + # Verify file operations + assert mock_file.call_count > 0 + + +def test_pull_missing_api_key(sample_df, mock_schema): + with patch("os.environ.get") as mock_env_get: + mock_env_get.return_value = None + with pytest.raises(PandaAIApiKeyError): + df = DataFrame(sample_df, path="test/path", schema=mock_schema) + df.pull() + + +def test_pull_api_error(mock_env, sample_df, mock_schema): + with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session: + mock_response = Mock() + mock_response.status_code = 404 + mock_session.return_value.get.return_value = mock_response + + df = DataFrame(sample_df, path="test/path", schema=mock_schema) + with pytest.raises(DatasetNotFound, match="Remote dataset not found to pull!"): + df.pull() + + +def test_pull_file_exists(mock_env, sample_df, mock_zip_content, mock_schema, tmp_path): + with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, \ + patch("pandasai.dataframe.base.find_project_root") as mock_root, \ + patch("pandasai.DatasetLoader.create_loader_from_path") as mock_loader, \ + patch("builtins.open", mock_open()) as mock_file, \ + patch("os.path.exists") as mock_exists, \ + patch("os.makedirs") as mock_makedirs: + + # Setup mocks + mock_response = Mock() + mock_response.status_code = 200 + mock_response.content = mock_zip_content + mock_session.return_value.get.return_value = mock_response + mock_root.return_value = str(tmp_path) + mock_exists.return_value = True + + mock_loader_instance = Mock() + mock_loader_instance.load.return_value = DataFrame(sample_df, schema=mock_schema) + mock_loader.return_value = mock_loader_instance + + # Create DataFrame instance and call pull + df = DataFrame(sample_df, path="test/path", schema=mock_schema) + df.pull() + + # Verify directory creation + mock_makedirs.assert_called_with(os.path.dirname(os.path.join(str(tmp_path), "datasets", "test/path", "test.csv")), exist_ok=True) From 27bab00e826c7f9eddd63a6b9e56b5659b1595df Mon Sep 17 00:00:00 2001 From: Gabriele Venturi Date: Fri, 31 Jan 2025 17:22:13 +0100 Subject: [PATCH 6/6] lint: fix linte --- tests/test_memory.py | 35 ++++++------- tests/unit_tests/dataframe/test_pull.py | 63 +++++++++++++++--------- tests/unit_tests/helpers/test_logger.py | 35 ++++++++++--- tests/unit_tests/helpers/test_session.py | 6 +-- 4 files changed, 85 insertions(+), 54 deletions(-) diff --git a/tests/test_memory.py b/tests/test_memory.py index 15d3ebec9..907a6d9c8 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -1,4 +1,3 @@ -import pytest from pandasai.helpers.memory import Memory @@ -9,36 +8,32 @@ def test_to_json_empty_memory(): def test_to_json_with_messages(): memory = Memory() - + # Add test messages memory.add("Hello", is_user=True) memory.add("Hi there!", is_user=False) memory.add("How are you?", is_user=True) - + expected_json = [ {"role": "user", "message": "Hello"}, {"role": "assistant", "message": "Hi there!"}, - {"role": "user", "message": "How are you?"} + {"role": "user", "message": "How are you?"}, ] - + assert memory.to_json() == expected_json def test_to_json_message_order(): memory = Memory() - + # Add messages in specific order - messages = [ - ("Message 1", True), - ("Message 2", False), - ("Message 3", True) - ] - + messages = [("Message 1", True), ("Message 2", False), ("Message 3", True)] + for msg, is_user in messages: memory.add(msg, is_user=is_user) - + result = memory.to_json() - + # Verify order is preserved assert len(result) == 3 assert result[0]["message"] == "Message 1" @@ -55,13 +50,13 @@ def test_to_openai_messages_with_agent_description(): memory = Memory(agent_description="I am a helpful assistant") memory.add("Hello", is_user=True) memory.add("Hi there!", is_user=False) - + expected_messages = [ {"role": "system", "content": "I am a helpful assistant"}, {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there!"} + {"role": "assistant", "content": "Hi there!"}, ] - + assert memory.to_openai_messages() == expected_messages @@ -70,11 +65,11 @@ def test_to_openai_messages_without_agent_description(): memory.add("Hello", is_user=True) memory.add("Hi there!", is_user=False) memory.add("How are you?", is_user=True) - + expected_messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, - {"role": "user", "content": "How are you?"} + {"role": "user", "content": "How are you?"}, ] - + assert memory.to_openai_messages() == expected_messages diff --git a/tests/unit_tests/dataframe/test_pull.py b/tests/unit_tests/dataframe/test_pull.py index d6da922fa..b0cfe6f2d 100644 --- a/tests/unit_tests/dataframe/test_pull.py +++ b/tests/unit_tests/dataframe/test_pull.py @@ -1,13 +1,18 @@ import os -import pytest -from unittest.mock import patch, Mock, mock_open from io import BytesIO +from unittest.mock import Mock, mock_open, patch from zipfile import ZipFile import pandas as pd +import pytest + +from pandasai.data_loader.semantic_layer_schema import ( + Column, + SemanticLayerSchema, + Source, +) from pandasai.dataframe.base import DataFrame -from pandasai.exceptions import PandaAIApiKeyError, DatasetNotFound -from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema, Column, Source +from pandasai.exceptions import DatasetNotFound, PandaAIApiKeyError @pytest.fixture @@ -41,20 +46,22 @@ def mock_schema(): def test_pull_success(mock_env, sample_df, mock_zip_content, mock_schema, tmp_path): - with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, \ - patch("pandasai.dataframe.base.find_project_root") as mock_root, \ - patch("pandasai.DatasetLoader.create_loader_from_path") as mock_loader, \ - patch("builtins.open", mock_open()) as mock_file: - + with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, patch( + "pandasai.dataframe.base.find_project_root" + ) as mock_root, patch( + "pandasai.DatasetLoader.create_loader_from_path" + ) as mock_loader, patch("builtins.open", mock_open()) as mock_file: # Setup mocks mock_response = Mock() mock_response.status_code = 200 mock_response.content = mock_zip_content mock_session.return_value.get.return_value = mock_response mock_root.return_value = str(tmp_path) - + mock_loader_instance = Mock() - mock_loader_instance.load.return_value = DataFrame(sample_df, schema=mock_schema) + mock_loader_instance.load.return_value = DataFrame( + sample_df, schema=mock_schema + ) mock_loader.return_value = mock_loader_instance # Create DataFrame instance and call pull @@ -64,8 +71,11 @@ def test_pull_success(mock_env, sample_df, mock_zip_content, mock_schema, tmp_pa # Verify API call mock_session.return_value.get.assert_called_once_with( "/datasets/pull", - headers={"accept": "application/json", "x-authorization": "Bearer test_api_key"}, - params={"path": "test/path"} + headers={ + "accept": "application/json", + "x-authorization": "Bearer test_api_key", + }, + params={"path": "test/path"}, ) # Verify file operations @@ -92,13 +102,13 @@ def test_pull_api_error(mock_env, sample_df, mock_schema): def test_pull_file_exists(mock_env, sample_df, mock_zip_content, mock_schema, tmp_path): - with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, \ - patch("pandasai.dataframe.base.find_project_root") as mock_root, \ - patch("pandasai.DatasetLoader.create_loader_from_path") as mock_loader, \ - patch("builtins.open", mock_open()) as mock_file, \ - patch("os.path.exists") as mock_exists, \ - patch("os.makedirs") as mock_makedirs: - + with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, patch( + "pandasai.dataframe.base.find_project_root" + ) as mock_root, patch( + "pandasai.DatasetLoader.create_loader_from_path" + ) as mock_loader, patch("builtins.open", mock_open()) as mock_file, patch( + "os.path.exists" + ) as mock_exists, patch("os.makedirs") as mock_makedirs: # Setup mocks mock_response = Mock() mock_response.status_code = 200 @@ -106,9 +116,11 @@ def test_pull_file_exists(mock_env, sample_df, mock_zip_content, mock_schema, tm mock_session.return_value.get.return_value = mock_response mock_root.return_value = str(tmp_path) mock_exists.return_value = True - + mock_loader_instance = Mock() - mock_loader_instance.load.return_value = DataFrame(sample_df, schema=mock_schema) + mock_loader_instance.load.return_value = DataFrame( + sample_df, schema=mock_schema + ) mock_loader.return_value = mock_loader_instance # Create DataFrame instance and call pull @@ -116,4 +128,9 @@ def test_pull_file_exists(mock_env, sample_df, mock_zip_content, mock_schema, tm df.pull() # Verify directory creation - mock_makedirs.assert_called_with(os.path.dirname(os.path.join(str(tmp_path), "datasets", "test/path", "test.csv")), exist_ok=True) + mock_makedirs.assert_called_with( + os.path.dirname( + os.path.join(str(tmp_path), "datasets", "test/path", "test.csv") + ), + exist_ok=True, + ) diff --git a/tests/unit_tests/helpers/test_logger.py b/tests/unit_tests/helpers/test_logger.py index a47a9a907..5ac68235a 100644 --- a/tests/unit_tests/helpers/test_logger.py +++ b/tests/unit_tests/helpers/test_logger.py @@ -1,4 +1,5 @@ import logging + from pandasai.helpers.logger import Logger @@ -6,26 +7,39 @@ def test_verbose_setter(): # Initialize logger with verbose=False logger = Logger(verbose=False) assert logger._verbose is False - assert not any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert not any( + isinstance(handler, logging.StreamHandler) + for handler in logger._logger.handlers + ) # Set verbose to True logger.verbose = True assert logger._verbose is True - assert any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert any( + isinstance(handler, logging.StreamHandler) + for handler in logger._logger.handlers + ) assert len(logger._logger.handlers) == 1 # Set verbose to False logger.verbose = False assert logger._verbose is False - assert not any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert not any( + isinstance(handler, logging.StreamHandler) + for handler in logger._logger.handlers + ) assert len(logger._logger.handlers) == 0 # Set verbose to True again to ensure multiple toggles work logger.verbose = True assert logger._verbose is True - assert any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers) + assert any( + isinstance(handler, logging.StreamHandler) + for handler in logger._logger.handlers + ) assert len(logger._logger.handlers) == 1 + def test_save_logs_property(): # Initialize logger with save_logs=False logger = Logger(save_logs=False, verbose=False) @@ -34,22 +48,27 @@ def test_save_logs_property(): # Enable save_logs logger.save_logs = True assert logger.save_logs is True - assert any(isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers) + assert any( + isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers + ) # Disable save_logs logger.save_logs = False assert logger.save_logs is False - assert not any(isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers) + assert not any( + isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers + ) + def test_save_logs_property(): # When logger is initialized with save_logs=True (default), it should have handlers logger = Logger(save_logs=True) assert logger.save_logs is True - + # When logger is initialized with save_logs=False, it should still have handlers if verbose=True logger = Logger(save_logs=False, verbose=True) assert logger.save_logs is True - + # When both save_logs and verbose are False, there should be no handlers logger = Logger(save_logs=False, verbose=False) logger._logger.handlers = [] # Reset handlers to match the property's expected behavior diff --git a/tests/unit_tests/helpers/test_session.py b/tests/unit_tests/helpers/test_session.py index 264766ab0..3b427c939 100644 --- a/tests/unit_tests/helpers/test_session.py +++ b/tests/unit_tests/helpers/test_session.py @@ -5,7 +5,7 @@ import requests from pandasai.constants import DEFAULT_API_URL -from pandasai.exceptions import PandaAIApiKeyError, PandaAIApiCallError +from pandasai.exceptions import PandaAIApiCallError, PandaAIApiKeyError from pandasai.helpers.session import Session, get_pandaai_session @@ -148,7 +148,7 @@ def test_make_request_error_response(mock_request): session = Session(api_key="test-key") with pytest.raises(PandaAIApiCallError) as exc_info: session.make_request("POST", "/test") - + assert str(exc_info.value) == "Bad request" @@ -161,7 +161,7 @@ def test_make_request_network_error(mock_request): session = Session(api_key="test-key") with pytest.raises(PandaAIApiCallError) as exc_info: session.make_request("GET", "/test") - + assert "Request failed: Network error" in str(exc_info.value)