fix(QueryTracker): publish query tracker results #773
Conversation
WalkthroughThe changes involve enhancing a data pipeline's code execution tracking and testing mechanisms. A new query tracker is introduced to publish updates post pipeline execution, and the Changes
TipsChat with CodeRabbit Bot (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Review Status
Actionable comments generated: 1
Configuration used: CodeRabbit UI
Files selected for processing (2)
- pandasai/smart_datalake/init.py (1 hunks)
- tests/test_smartdatalake.py (1 hunks)
Files skipped from review due to trivial changes (1)
- pandasai/smart_datalake/init.py
| "value": "There are 10 countries in the dataframe.", | ||
| } | ||
|
|
||
| @patch.object( | ||
| CodeManager, | ||
| "execute_code", | ||
| return_value={ | ||
| "type": "string", | ||
| "value": "There are 10 countries in the dataframe.", | ||
| }, | ||
| ) | ||
| @patch("pandasai.helpers.query_exec_tracker.QueryExecTracker.publish") | ||
| def test_query_tracker_publish_called_in_chat_method( | ||
| self, mock_query_tracker_publish, _mocked_method, smart_datalake: SmartDatalake | ||
| ): | ||
| assert smart_datalake.last_result is None | ||
|
|
||
| _mocked_method.__name__ = "execute_code" | ||
|
|
||
| smart_datalake.chat("How many countries are in the dataframe?") | ||
| mock_query_tracker_publish.assert_called() | ||
|
|
There was a problem hiding this comment.
The new test method test_query_tracker_publish_called_in_chat_method is correctly structured to test the invocation of the query_tracker_publish function. However, there are a few points to consider:
- The test does not verify the parameters with which
mock_query_tracker_publishis called. It is important to ensure that the correct arguments are being passed to the function. - The test assumes that
smart_datalake.last_resultisNonebefore thechatmethod is called, which is good. However, it does not check the state oflast_resultafter thechatmethod is called to ensure that it has been updated correctly. - The test uses patching correctly, but it's not clear if the
mock_query_tracker_publishis being reset after each test. This could potentially affect other tests if the mock state is not isolated.
Consider adding assertions to check the arguments passed to mock_query_tracker_publish and the state of last_result after the chat method is called. Also, ensure that mocks are reset after each test to maintain test isolation.
- smart_datalake.chat("How many countries are in the dataframe?")
+ question = "How many countries are in the dataframe?"
+ smart_datalake.chat(question)
+ mock_query_tracker_publish.assert_called_with(question, ANY)
+ assert smart_datalake.last_result is not NoneCommitable suggestion
[!IMPORTANT]
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| "value": "There are 10 countries in the dataframe.", | |
| } | |
| @patch.object( | |
| CodeManager, | |
| "execute_code", | |
| return_value={ | |
| "type": "string", | |
| "value": "There are 10 countries in the dataframe.", | |
| }, | |
| ) | |
| @patch("pandasai.helpers.query_exec_tracker.QueryExecTracker.publish") | |
| def test_query_tracker_publish_called_in_chat_method( | |
| self, mock_query_tracker_publish, _mocked_method, smart_datalake: SmartDatalake | |
| ): | |
| assert smart_datalake.last_result is None | |
| _mocked_method.__name__ = "execute_code" | |
| smart_datalake.chat("How many countries are in the dataframe?") | |
| mock_query_tracker_publish.assert_called() | |
| @patch.object( | |
| CodeManager, | |
| "execute_code", | |
| return_value={ | |
| "type": "string", | |
| "value": "There are 10 countries in the dataframe.", | |
| }, | |
| ) | |
| @patch("pandasai.helpers.query_exec_tracker.QueryExecTracker.publish") | |
| def test_query_tracker_publish_called_in_chat_method( | |
| self, mock_query_tracker_publish, _mocked_method, smart_datalake: SmartDatalake | |
| ): | |
| assert smart_datalake.last_result is None | |
| _mocked_method.__name__ = "execute_code" | |
| question = "How many countries are in the dataframe?" | |
| smart_datalake.chat(question) | |
| mock_query_tracker_publish.assert_called_with(question, ANY) | |
| assert smart_datalake.last_result is not None |
Codecov ReportAll modified and coverable lines are covered by tests ✅
❗ Your organization needs to install the Codecov GitHub app to enable full functionality. Additional details and impacted files@@ Coverage Diff @@
## main #773 +/- ##
=======================================
Coverage 85.09% 85.09%
=======================================
Files 88 88
Lines 3824 3825 +1
=======================================
+ Hits 3254 3255 +1
Misses 570 570 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
Review Status
Actionable comments generated: 1
Configuration used: CodeRabbit UI
Files selected for processing (2)
- pandasai/pipelines/smart_datalake_chat/code_execution.py (1 hunks)
- tests/pipelines/smart_datalake/test_code_execution.py (2 hunks)
Additional comments: 2
pandasai/pipelines/smart_datalake_chat/code_execution.py (1)
- 42-52: The change to use
pipeline_context.query_exec_tracker.execute_funcfor executing code is a significant improvement in terms of maintainability and encapsulation. It centralizes the execution logic and potentially allows for better tracking and management of code execution within the pipeline context. Ensure that theexecute_funcmethod is thoroughly tested, especially for its error handling and retry logic.tests/pipelines/smart_datalake/test_code_execution.py (1)
- 100-107: The
context.get_intermediate_valuemethod is mocked to return different objects based on the key provided. However, the mock forcontext._query_exec_trackerandcontext.query_exec_tracker.execute_funcis set up but not used in thetest_code_execution_successful_with_no_exceptionsmethod. Ensure that the mocks are being utilized correctly in the test case to reflect the intended behavior.
| raise Exception("Unit test exception") | ||
| return "Mocked Result after retry" | ||
|
|
||
| # Conditional return of execute_func method based arguments it is called with | ||
| def mock_execute_func(*args, **kwargs): | ||
| if isinstance(args[0], Mock) and args[0].name == "execute_code": | ||
| return mock_execute_code(*args, **kwargs) | ||
| else: | ||
| return [ | ||
| "Interuppted Code", | ||
| "Exception Testing", | ||
| "Successful after Retry", | ||
| ] | ||
|
|
||
| mock_code_manager = Mock() | ||
| mock_code_manager.execute_code = Mock(side_effect=mock_execute_code) | ||
| mock_code_manager.execute_code = Mock() | ||
| mock_code_manager.execute_code.name = "execute_code" | ||
|
|
||
| context._query_exec_tracker = Mock() | ||
| context.query_exec_tracker.execute_func = Mock( | ||
| return_value=[ | ||
| "Interuppted Code", | ||
| "Exception Testing", | ||
| "Successful after Retry", | ||
| ] | ||
| ) | ||
|
|
||
| context.query_exec_tracker.execute_func = Mock(side_effect=mock_execute_func) | ||
|
|
||
| def mock_intermediate_values(key: str): |
There was a problem hiding this comment.
The mock_execute_func is designed to simulate different behaviors based on the arguments it receives. However, the conditional logic within mock_execute_func seems to be based on the name attribute of the Mock object, which is not a standard attribute. This could lead to unexpected behavior or test failures if the name attribute is not set elsewhere in the code. Consider using a more reliable condition to differentiate the behavior of mock_execute_func.
Summary by CodeRabbit
New Features
Tests
chatmethod execution.