diff --git a/codecarbon/core/api_client.py b/codecarbon/core/api_client.py index 58f4932cb..1d8675f7b 100644 --- a/codecarbon/core/api_client.py +++ b/codecarbon/core/api_client.py @@ -70,6 +70,14 @@ def _get_headers(self): headers["Authorization"] = f"Bearer {self.access_token}" return headers + def _request(self, method, url, payload=None, expected_status=200): + headers = self._get_headers() + response = method(url=url, json=payload, timeout=2, headers=headers) + if response.status_code != expected_status: + self._log_error(url, payload or {}, response) + response.raise_for_status() + return response + def set_access_token(self, token: str): """This method sets the access token to be used for the API. Args: @@ -82,32 +90,20 @@ def check_auth(self): Check API access to user account """ url = self.url + "/auth/check" - headers = self._get_headers() - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return None - return r.json() + return self._request(requests.get, url).json() def get_list_organizations(self): """ List all organizations """ url = self.url + "/organizations" - headers = self._get_headers() - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return None - return r.json() + return self._request(requests.get, url).json() def check_organization_exists(self, organization_name: str): """ Check if an organization exists """ organizations = self.get_list_organizations() - if organizations is None: - return False for organization in organizations: if organization["name"] == organization_name: return organization @@ -125,49 +121,31 @@ def create_organization(self, organization: OrganizationCreate): ) return organization else: - headers = self._get_headers() - r = requests.post(url=url, json=payload, timeout=2, headers=headers) - if r.status_code != 201: - self._log_error(url, payload, r) - return None - return r.json() + return self._request( + requests.post, url, payload=payload, expected_status=201 + ).json() def get_organization(self, organization_id): """ Get an organization """ - headers = self._get_headers() url = self.url + "/organizations/" + organization_id - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return None - return r.json() + return self._request(requests.get, url).json() def update_organization(self, organization: OrganizationCreate): """ Update an organization """ payload = dataclasses.asdict(organization) - headers = self._get_headers() url = self.url + "/organizations/" + organization.id - r = requests.patch(url=url, json=payload, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, payload, r) - return None - return r.json() + return self._request(requests.patch, url, payload=payload).json() def list_projects_from_organization(self, organization_id): """ List all projects """ url = self.url + "/organizations/" + organization_id + "/projects" - headers = self._get_headers() - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return None - return r.json() + return self._request(requests.get, url).json() def create_project(self, project: ProjectCreate): """ @@ -175,24 +153,16 @@ def create_project(self, project: ProjectCreate): """ payload = dataclasses.asdict(project) url = self.url + "/projects" - headers = self._get_headers() - r = requests.post(url=url, json=payload, timeout=2, headers=headers) - if r.status_code != 201: - self._log_error(url, payload, r) - return None - return r.json() + return self._request( + requests.post, url, payload=payload, expected_status=201 + ).json() def get_project(self, project_id): """ Get a project """ url = self.url + "/projects/" + project_id - headers = self._get_headers() - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return None - return r.json() + return self._request(requests.get, url).json() def add_emission(self, carbon_emission: dict): assert self.experiment_id is not None @@ -230,18 +200,10 @@ def add_emission(self, carbon_emission: dict): ram_utilization_percent=carbon_emission.get("ram_utilization_percent"), wue=carbon_emission.get("wue", 0), ) - try: - payload = dataclasses.asdict(emission) - url = self.url + "/emissions" - headers = self._get_headers() - r = requests.post(url=url, json=payload, timeout=2, headers=headers) - if r.status_code != 201: - self._log_error(url, payload, r) - return False - logger.debug(f"ApiClient - Successful upload emission {payload} to {url}") - except Exception as e: - logger.error(e, exc_info=True) - return False + payload = dataclasses.asdict(emission) + url = self.url + "/emissions" + self._request(requests.post, url, payload=payload, expected_status=201) + logger.debug(f"ApiClient - Successful upload emission {payload} to {url}") return True def _create_run(self, experiment_id: str): @@ -275,11 +237,7 @@ def _create_run(self, experiment_id: str): ) payload = dataclasses.asdict(run) url = self.url + "/runs" - headers = self._get_headers() - r = requests.post(url=url, json=payload, timeout=2, headers=headers) - if r.status_code != 201: - self._log_error(url, payload, r) - return None + r = self._request(requests.post, url, payload=payload, expected_status=201) self.run_id = r.json()["id"] logger.info( "ApiClient Successfully registered your run on the API.\n\n" @@ -292,6 +250,8 @@ def _create_run(self, experiment_id: str): f"Failed to connect to API, please check the configuration. {e}", exc_info=False, ) + except requests.exceptions.HTTPError: + raise except Exception as e: logger.error(e, exc_info=True) @@ -300,12 +260,7 @@ def list_experiments_from_project(self, project_id: str): List all experiments for a project """ url = self.url + "/projects/" + project_id + "/experiments" - headers = self._get_headers() - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return [] - return r.json() + return self._request(requests.get, url).json() def set_experiment(self, experiment_id: str): """ @@ -320,24 +275,16 @@ def add_experiment(self, experiment: ExperimentCreate): """ payload = dataclasses.asdict(experiment) url = self.url + "/experiments" - headers = self._get_headers() - r = requests.post(url=url, json=payload, timeout=2, headers=headers) - if r.status_code != 201: - self._log_error(url, payload, r) - return None - return r.json() + return self._request( + requests.post, url, payload=payload, expected_status=201 + ).json() def get_experiment(self, experiment_id): """ Get an experiment by id """ url = self.url + "/experiments/" + experiment_id - headers = self._get_headers() - r = requests.get(url=url, timeout=2, headers=headers) - if r.status_code != 200: - self._log_error(url, {}, r) - return None - return r.json() + return self._request(requests.get, url).json() def _log_error(self, url, payload, response): if len(payload) > 0: diff --git a/tests/test_api_call.py b/tests/test_api_call.py index 39822ece7..bb782ad39 100644 --- a/tests/test_api_call.py +++ b/tests/test_api_call.py @@ -2,6 +2,7 @@ import unittest from uuid import uuid4 +import requests import requests_mock from codecarbon.core.api_client import ApiClient @@ -146,7 +147,8 @@ def test_check_auth_returns_none_on_error(self): create_run_automatically=False, ) - self.assertIsNone(api.check_auth()) + with self.assertRaises(requests.exceptions.HTTPError): + api.check_auth() def test_check_organization_exists_returns_false_when_list_fails(self): with requests_mock.Mocker() as m: @@ -156,7 +158,8 @@ def test_check_organization_exists_returns_false_when_list_fails(self): create_run_automatically=False, ) - self.assertFalse(api.check_organization_exists("missing")) + with self.assertRaises(requests.exceptions.HTTPError): + api.check_organization_exists("missing") def test_create_organization_skips_when_name_exists(self): organization = OrganizationCreate(name="existing", description="desc") @@ -236,7 +239,7 @@ def test_add_emission_returns_false_on_unsuccessful_post(self): ) api.run_id = "run-1" - self.assertFalse( + with self.assertRaises(requests.exceptions.HTTPError): api.add_emission( { "duration": 2, @@ -251,7 +254,6 @@ def test_add_emission_returns_false_on_unsuccessful_post(self): "energy_consumed": 0.2, } ) - ) def test_create_run_returns_none_on_unsuccessful_status(self): with requests_mock.Mocker() as m: @@ -264,7 +266,8 @@ def test_create_run_returns_none_on_unsuccessful_status(self): create_run_automatically=False, ) - self.assertIsNone(api._create_run("experiment_id")) + with self.assertRaises(requests.exceptions.HTTPError): + api._create_run("experiment_id") self.assertIsNone(api.run_id) def test_list_experiments_from_project_returns_empty_list_on_error(self): @@ -279,7 +282,8 @@ def test_list_experiments_from_project_returns_empty_list_on_error(self): create_run_automatically=False, ) - self.assertEqual(api.list_experiments_from_project("proj-1"), []) + with self.assertRaises(requests.exceptions.HTTPError): + api.list_experiments_from_project("proj-1") def test_set_experiment_updates_value(self): api = ApiClient(endpoint_url="http://test.com", create_run_automatically=False) @@ -303,7 +307,8 @@ def test_add_experiment_returns_none_on_error(self): create_run_automatically=False, ) - self.assertIsNone(api.add_experiment(experiment)) + with self.assertRaises(requests.exceptions.HTTPError): + api.add_experiment(experiment) def test_get_experiment_returns_none_on_error(self): with requests_mock.Mocker() as m: @@ -313,4 +318,5 @@ def test_get_experiment_returns_none_on_error(self): create_run_automatically=False, ) - self.assertIsNone(api.get_experiment("exp-1")) + with self.assertRaises(requests.exceptions.HTTPError): + api.get_experiment("exp-1")