refactor: better error handling

ThreatX · May 2, 2024 · 959cf4e · 959cf4e
1 parent 238e843
commit 959cf4e
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/threatx_api_client/__init__.py b/threatx_api_client/__init__.py
@@ -6,7 +6,7 @@
     TXAPIIncorrectCommandError,
     TXAPIIncorrectEnvironmentError,
     TXAPIIncorrectTokenError,
-    TXAPIResponseError,
+    TXAPIResponseError, TXAPIError,
 )
 
 
@@ -67,9 +67,11 @@ async def __post(self, session, path: str, post_payload: dict):
                     self.session_token = self.__get_session_token()
                     return self.__post(session, path, post_payload)
                 elif response_error_data:
-                    raise TXAPIResponseError(response_error_data)
+                    error_msg = {marker_var: response_error_data} if marker_var else response_error_data
+                    raise TXAPIResponseError(error_msg)
                 else:
-                    pass
+                    error_msg = {marker_var: raw_response} if marker_var else raw_response
+                    raise TXAPIError(error_msg)
 
     async def __process_response(self, path: str, available_commands: list, payloads):
         if isinstance(payloads, dict):
@@ -85,7 +87,7 @@ async def __process_response(self, path: str, available_commands: list, payloads
                     session,
                     path,
                     {"token": self.session_token, **payload}) for payload in payloads
-            ))
+            ), return_exceptions=True)
 
         if len(responses) == 1:
             return responses[0]