diff --git a/src/python/gem5/resources/client_api/atlasclient.py b/src/python/gem5/resources/client_api/atlasclient.py index a39f478458..511ef71fed 100644 --- a/src/python/gem5/resources/client_api/atlasclient.py +++ b/src/python/gem5/resources/client_api/atlasclient.py @@ -30,7 +30,7 @@ import json import time import itertools from .abstract_client import AbstractClient - +from urllib.error import HTTPError from m5.util import warn @@ -77,7 +77,7 @@ class AtlasClient(AbstractClient): def get_token(self): return self._atlas_http_json_req( - self.url, + self.authUrl, data_json={"key": self.apiKey}, headers={"Content-Type": "application/json"}, purpose_of_request="Get Access Token with API key", @@ -122,7 +122,7 @@ class AtlasClient(AbstractClient): try: response = request.urlopen(req) break - except Exception as e: + except HTTPError as e: if attempt >= max_failed_attempts: raise AtlasClientHttpJsonRequestError( client=self, @@ -167,7 +167,10 @@ class AtlasClient(AbstractClient): } resources = self._atlas_http_json_req( - url, data=data, headers=headers, purpose_of_request="Get Resources" + url, + data_json=data, + headers=headers, + purpose_of_request="Get Resources", )["documents"] # I do this as a lazy post-processing step because I can't figure out diff --git a/tests/pyunit/stdlib/resources/pyunit_client_wrapper_checks.py b/tests/pyunit/stdlib/resources/pyunit_client_wrapper_checks.py index f190b1ed5f..66b934a16f 100644 --- a/tests/pyunit/stdlib/resources/pyunit_client_wrapper_checks.py +++ b/tests/pyunit/stdlib/resources/pyunit_client_wrapper_checks.py @@ -34,6 +34,10 @@ import io import contextlib from pathlib import Path +from gem5.resources.client_api.atlasclient import ( + AtlasClientHttpJsonRequestError, +) + mock_json_path = Path(__file__).parent / "refs/resources.json" mock_config_json = { "sources": { @@ -419,21 +423,11 @@ class ClientWrapperTestSuite(unittest.TestCase): @patch("urllib.request.urlopen", side_effect=mocked_requests_post) def test_invalid_auth_url(self, mock_get): resource_id = "test-resource" - f = io.StringIO() - with self.assertRaises(Exception) as context: - with contextlib.redirect_stderr(f): - get_resource_json_obj( - resource_id, - gem5_version="develop", - ) - self.assertTrue( - "Error getting resources from client gem5-resources:" - " Panic: Not found" in str(f.getvalue()) - ) - self.assertTrue( - "Resource with ID 'test-resource' not found." - in str(context.exception) - ) + with self.assertRaises(AtlasClientHttpJsonRequestError) as context: + get_resource_json_obj( + resource_id, + gem5_version="develop", + ) @patch( "gem5.resources.client.clientwrapper", @@ -442,21 +436,11 @@ class ClientWrapperTestSuite(unittest.TestCase): @patch("urllib.request.urlopen", side_effect=mocked_requests_post) def test_invalid_url(self, mock_get): resource_id = "test-resource" - f = io.StringIO() - with self.assertRaises(Exception) as context: - with contextlib.redirect_stderr(f): - get_resource_json_obj( - resource_id, - gem5_version="develop", - ) - self.assertTrue( - "Error getting resources from client gem5-resources:" - " Panic: Not found" in str(f.getvalue()) - ) - self.assertTrue( - "Resource with ID 'test-resource' not found." - in str(context.exception) - ) + with self.assertRaises(AtlasClientHttpJsonRequestError) as context: + get_resource_json_obj( + resource_id, + gem5_version="develop", + ) @patch( "gem5.resources.client.clientwrapper", @@ -465,18 +449,8 @@ class ClientWrapperTestSuite(unittest.TestCase): @patch("urllib.request.urlopen", side_effect=mocked_requests_post) def test_invalid_url(self, mock_get): resource_id = "test-too-many" - f = io.StringIO() - with self.assertRaises(Exception) as context: - with contextlib.redirect_stderr(f): - get_resource_json_obj( - resource_id, - gem5_version="develop", - ) - self.assertTrue( - "Error getting resources from client gem5-resources:" - " Panic: Too many requests" in str(f.getvalue()) - ) - self.assertTrue( - "Resource with ID 'test-too-many' not found." - in str(context.exception) - ) + with self.assertRaises(AtlasClientHttpJsonRequestError) as context: + get_resource_json_obj( + resource_id, + gem5_version="develop", + )