diff --git a/jupiterone/client.py b/jupiterone/client.py index a9be725..422d24e 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 account id and API 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..708d573 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 account id and API token.' in str(ex.value) + + @responses.activate def test_tree_query_v1():