Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 32 additions & 85 deletions codecarbon/core/api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this helper!

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Indeed it elegant. I will do this in #1277 and will credit you @PSR94 in the release note.

I prefer to merge #1277 as it better address the problem as a whole.

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:
Expand All @@ -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
Expand All @@ -125,74 +121,48 @@ 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):
"""
Create a project
"""
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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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"
Expand All @@ -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)

Expand All @@ -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):
"""
Expand All @@ -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:
Expand Down
22 changes: 14 additions & 8 deletions tests/test_api_call.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import unittest
from uuid import uuid4

import requests
import requests_mock

from codecarbon.core.api_client import ApiClient
Expand Down Expand Up @@ -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:
Expand All @@ -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")
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand All @@ -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):
Expand All @@ -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)
Expand All @@ -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:
Expand All @@ -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")
Loading