Merge pull request #39 from mbengit/master

Fix ariel delete/cancel issues
IBM · Sep 11, 2020 · 852067b · 852067b
2 parents 1680072 + 774aec5
commit 852067b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 25 deletions.
diff --git a/qpylib/ariel.py b/qpylib/ariel.py
@@ -112,22 +112,21 @@ def delete(self, search_id, api_version='latest'):
         ''' Deletes a previous Ariel search.
               search_id: Ariel search ID
               api_version: QRadar API version to use, defaults to latest.
-            Returns a tuple containing search status and search ID.
+            Returns search status.
             Raises ArielError if the search could not be deleted.
         '''
         response = qpylib.REST('DELETE', ArielSearch.SEARCH_ENDPOINT.format(search_id),
                                headers=self._build_headers(api_version))
-        if response.status_code != 200:
+        if response.status_code != 202:
             raise ArielError('Ariel search {0} could not be deleted: HTTP {1} was returned'
                              .format(search_id, response.status_code))
-        response_json = response.json()
-        return (response_json['status'], response_json['search_id'])
+        return response.json()['status']
 
     def cancel(self, search_id, api_version='latest'):
         ''' Cancels an ongoing Ariel search.
               search_id: Ariel search ID.
               api_version: QRadar API version to use, defaults to latest.
-            Returns a tuple containing search status and search ID.
+            Returns search status.
             Raises ArielError if the search could not be cancelled.
         '''
         response = qpylib.REST('POST', ArielSearch.SEARCH_ENDPOINT.format(search_id),
@@ -136,8 +135,7 @@ def cancel(self, search_id, api_version='latest'):
         if response.status_code != 200:
             raise ArielError('Ariel search {0} could not be cancelled: HTTP {1} was returned'
                              .format(search_id, response.status_code))
-        response_json = response.json()
-        return (response_json['status'], response_json['search_id'])
+        return response.json()['status']
 
     @staticmethod
     def _build_headers(api_version):

diff --git a/test/test_ariel.py b/test/test_ariel.py
@@ -158,34 +158,26 @@ def test_results_range_start_and_end():
 
 @responses.activate
 def test_search_delete_failure():
-    responses.add('DELETE', DELETE_SEARCH, status=404,
-                  json={})
-    with pytest.raises(ArielError, match='Ariel search {0} could not be deleted'
+    responses.add('DELETE', DELETE_SEARCH, status=500)
+    with pytest.raises(ArielError,
+                       match='Ariel search {0} could not be deleted: HTTP 500 was returned'
                        .format(SEARCH_ID)):
         ArielSearch().delete(SEARCH_ID)
 
 @responses.activate
 def test_search_delete_success():
-    responses.add('DELETE', DELETE_SEARCH, status=200,
-                  json={'status': 'COMPLETED', 'search_id': SEARCH_ID})
-    status, search_id = ArielSearch().delete(SEARCH_ID)
-
-    assert status == 'COMPLETED'
-    assert search_id == SEARCH_ID
+    responses.add('DELETE', DELETE_SEARCH, status=202, json={'status': 'COMPLETED'})
+    assert ArielSearch().delete(SEARCH_ID) == 'COMPLETED'
 
 @responses.activate
 def test_search_cancel_failure():
-    responses.add('POST', CANCEL_SEARCH, status=404,
-                  json={})
-    with pytest.raises(ArielError, match='Ariel search {0} could not be cancelled'
+    responses.add('POST', CANCEL_SEARCH, status=500)
+    with pytest.raises(ArielError,
+                       match='Ariel search {0} could not be cancelled: HTTP 500 was returned'
                        .format(SEARCH_ID)):
         ArielSearch().cancel(SEARCH_ID)
 
 @responses.activate
 def test_search_cancel_success():
-    responses.add('POST', CANCEL_SEARCH, status=200,
-                  json={'status': 'COMPLETED', 'search_id': SEARCH_ID})
-    status, search_id = ArielSearch().cancel(SEARCH_ID)
-
-    assert status == 'COMPLETED'
-    assert search_id == SEARCH_ID
+    responses.add('POST', CANCEL_SEARCH, status=200, json={'status': 'COMPLETED'})
+    assert ArielSearch().cancel(SEARCH_ID) == 'COMPLETED'