Cybereason driver query return instance name in dataframe (#804)

* Add max retry parameter to the execution * Add tests and raise exception after max retry * doc and keyword params * return cybereason instance name in dataframe --------- Co-authored-by: Meroujan.Antonyan <[email protected]> Co-authored-by: Ian Hellen <[email protected]>
microsoft · Nov 23, 2024 · eacda5c · eacda5c
1 parent 6db4cc3
commit eacda5c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/msticpy/data/drivers/cybereason_driver.py b/msticpy/data/drivers/cybereason_driver.py
@@ -159,7 +159,7 @@ def query(
             )
         else:
             df_result = self._format_result_to_dataframe(result=response)
-
+        df_result["instance"] = self.instance
         return df_result
 
     def _exec_paginated_queries(

diff --git a/tests/data/drivers/test_cybereason_driver.py b/tests/data/drivers/test_cybereason_driver.py
@@ -281,6 +281,7 @@ def test_query(driver):
         check.is_true(connect.called or driver.connected)
         check.is_true(query.called)
         check.is_instance(data, pd.DataFrame)
+        check.is_true("instance" in data.columns)
 
 @respx.mock
 def test_partial_success_query(driver):