fix(sql): update comment and test case

sinaptik-ai · Jan 31, 2025 · 8b29e2c · 8b29e2c
1 parent 51dac82
commit 8b29e2c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/pandasai/data_loader/sql_loader.py b/pandasai/data_loader/sql_loader.py
@@ -51,7 +51,7 @@ def execute_query(self, query: str, params: Optional[list] = None) -> pd.DataFra
 
         except ModuleNotFoundError as e:
             raise ImportError(
-                f"{source_type.capitalize()} connector not found. Please install the pandasai_sql[{source_type}] library.e.g. `pip install pandasai_sql[{source_type}]`"
+                f"{source_type.capitalize()} connector not found. Please install the pandasai_sql[{source_type}] library, e.g. `pip install pandasai_sql[{source_type}]`."
             ) from e
 
         except Exception as e:

diff --git a/tests/unit_tests/data_loader/test_sql_loader.py b/tests/unit_tests/data_loader/test_sql_loader.py
@@ -213,9 +213,12 @@ def test_mysql_malicious_with_no_import(self, mysql_schema):
                 )
             )
             mocked_exec_function.return_value = mock_df
-            mock_loader_function.side_effect = ModuleNotFoundError("Error")
+
+            mock_exec_function = MagicMock()
+            mock_loader_function.return_value = mock_exec_function
+            mock_exec_function.side_effect = ModuleNotFoundError("Error")
             loader = SQLDatasetLoader(mysql_schema, "test/users")
-            mock_sql_query.return_value = False
+            mock_sql_query.return_value = True
             logging.debug("Loading schema from dataset path: %s", loader)
             with pytest.raises(ImportError):
                 loader.execute_query("select * from users")