From d75d4954b83338ce14ffeddf075ea879d7a97b22 Mon Sep 17 00:00:00 2001
From: Samuel Poulton <samuel.poulton@jupiterone.com>
Date: Thu, 1 Jun 2023 14:18:18 -0600
Subject: [PATCH 1/2] Update client to handle unknown error formats and 401s

---
 jupiterone/client.py | 11 +++++++++--
 tests/test_query.py  | 27 +++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/jupiterone/client.py b/jupiterone/client.py
index a9be725..9747719 100644
--- a/jupiterone/client.py
+++ b/jupiterone/client.py
@@ -105,12 +105,19 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict:
                     raise JupiterOneApiError(content.get('errors'))
                 return response.json()
 
+        elif response.status_code == 401:
+            raise JupiterOneApiError('401: Unauthorized. Please supply a valid token')
+
         elif response.status_code in [429, 500]:
             raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded')
 
         else:
-            content = json.loads(response._content)
-            raise JupiterOneApiError('{}:{}'.format(response.status_code, content.get('error')))
+            try:
+                content = json.loads(response._content)
+                raise JupiterOneApiError('{}: {}'.format(response.status_code, content.get('error') or 'Unknown Error'))
+            except ValueError as e:
+                raise JupiterOneApiError('{}: {}'.format(response.status_code, 'Unknown Error'));
+
 
     def query_v1(self, query: str, **kwargs) -> Dict:
         """ Performs a V1 graph query
diff --git a/tests/test_query.py b/tests/test_query.py
index 9228d7f..4444b45 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -95,6 +95,33 @@ def test_query_v1():
     assert response['data'][0]['entity']['_id'] == '1'
 
 
+@responses.activate
+def test_401():
+    with pytest.raises(Exception) as ex:
+
+        def response_401_callback(r):
+            headers = {
+                'Content-Type': 'application/json'
+            }
+
+            response = {
+                'test': ['Unauthorized']
+            }
+            return (401, headers, json.dumps(response))
+
+        responses.add_callback(
+            responses.POST, 'https://api.us.jupiterone.io/graphql',
+            callback=response_401_callback,
+            content_type='application/text',
+        )
+
+        j1 = JupiterOneClient(account='testAccount', token='testToken')
+        query = "find Host with _id='1'"
+        j1.query_v1(query)
+        
+    assert '401: Unauthorized. Please supply a valid token' in str(ex.value)
+
+
 @responses.activate
 def test_tree_query_v1():
 

From f3c15d3674c974138e9d04890548cde115bef9a9 Mon Sep 17 00:00:00 2001
From: Samuel Poulton <samuel.poulton@jupiterone.com>
Date: Thu, 1 Jun 2023 14:28:38 -0600
Subject: [PATCH 2/2] update error msg to match UI

---
 jupiterone/client.py | 2 +-
 tests/test_query.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/jupiterone/client.py b/jupiterone/client.py
index 9747719..422d24e 100644
--- a/jupiterone/client.py
+++ b/jupiterone/client.py
@@ -106,7 +106,7 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict:
                 return response.json()
 
         elif response.status_code == 401:
-            raise JupiterOneApiError('401: Unauthorized. Please supply a valid token')
+            raise JupiterOneApiError('401: Unauthorized. Please supply a valid account id and API token.')
 
         elif response.status_code in [429, 500]:
             raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded')
diff --git a/tests/test_query.py b/tests/test_query.py
index 4444b45..708d573 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -119,7 +119,7 @@ def response_401_callback(r):
         query = "find Host with _id='1'"
         j1.query_v1(query)
         
-    assert '401: Unauthorized. Please supply a valid token' in str(ex.value)
+    assert '401: Unauthorized. Please supply a valid account id and API token.' in str(ex.value)
 
 
 @responses.activate