Skip to content

Commit e8dc95e

Browse files
committed
Add unit tests for logfire and http output methods.
This requires adding logfire as a dev dependency so we can import it in tests. Also make it clearer which `EmissionsData` params are unnecessary for the `out` and `live_out` methods.
1 parent f44672f commit e8dc95e

6 files changed

Lines changed: 277 additions & 6 deletions

File tree

codecarbon/output_methods/http.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ class HTTPOutput(BaseOutput):
1919
def __init__(self, endpoint_url: str):
2020
self.endpoint_url: str = endpoint_url
2121

22-
def out(self, total: EmissionsData, delta: EmissionsData):
22+
def out(self, total: EmissionsData, _: EmissionsData):
2323
try:
2424
payload = dataclasses.asdict(total)
2525
payload["user"] = getpass.getuser()
@@ -56,14 +56,14 @@ def __init__(
5656
)
5757
self.run_id = self.api.run_id
5858

59-
def live_out(self, total: EmissionsData, delta: EmissionsData):
59+
def live_out(self, _, delta: EmissionsData):
6060
# Called at regular intervals
6161
try:
6262
self.api.add_emission(dataclasses.asdict(delta))
6363
except Exception as e:
6464
logger.error(e, exc_info=True)
6565

66-
def out(self, total: EmissionsData, delta: EmissionsData):
66+
def out(self, _, delta: EmissionsData):
6767
# Called on exit
6868
try:
6969
self.api.add_emission(dataclasses.asdict(delta))

codecarbon/output_methods/metrics/logfire.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def __init__(self):
5959
"codecarbon_ram_energy", unit="(kWh)", description="Energy used per RAM"
6060
)
6161

62-
def out(self, total: EmissionsData, delta: EmissionsData):
62+
def out(self, _, delta: EmissionsData):
6363
try:
6464
self.duration.add(delta.duration)
6565
self.emissions.add(delta.emissions)
@@ -75,5 +75,5 @@ def out(self, total: EmissionsData, delta: EmissionsData):
7575
except Exception as e:
7676
logger.error(e, exc_info=True)
7777

78-
def live_out(self, total: EmissionsData, delta: EmissionsData):
79-
self.out(total, delta)
78+
def live_out(self, _: EmissionsData, delta: EmissionsData):
79+
self.out(None, delta)

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ dev = [
9090
"requests",
9191
"requests-mock",
9292
"responses",
93+
"logfire>=1.0.1",
9394
]
9495
doc = [
9596
"sphinx",

tests/output_methods/test_http.py

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
import unittest
2+
from unittest.mock import MagicMock, patch
3+
4+
from codecarbon.output_methods.emissions_data import EmissionsData
5+
from codecarbon.output_methods.http import CodeCarbonAPIOutput, HTTPOutput
6+
7+
8+
class TestHTTPOutput(unittest.TestCase):
9+
def setUp(self):
10+
self.emissions_data = EmissionsData(
11+
timestamp="2023-01-01T00:00:00",
12+
project_name="test_project",
13+
run_id="test_run_id",
14+
experiment_id="test_experiment_id",
15+
duration=10,
16+
emissions=0.5,
17+
emissions_rate=0.05,
18+
cpu_power=20,
19+
gpu_power=30,
20+
ram_power=5,
21+
cpu_energy=200,
22+
gpu_energy=300,
23+
ram_energy=50,
24+
energy_consumed=550,
25+
water_consumed=0.1,
26+
country_name="Testland",
27+
country_iso_code="TS",
28+
region="Test Region",
29+
cloud_provider="Test Cloud",
30+
cloud_region="test-cloud-1",
31+
os="TestOS",
32+
python_version="3.8",
33+
codecarbon_version="2.0",
34+
cpu_count=4,
35+
cpu_model="Test CPU",
36+
gpu_count=1,
37+
gpu_model="Test GPU",
38+
longitude=0,
39+
latitude=0,
40+
ram_total_size=16,
41+
tracking_mode="machine",
42+
on_cloud="true",
43+
pue=1.5,
44+
wue=0.5,
45+
)
46+
self.url = "http://test.com/emissions"
47+
self.http_output = HTTPOutput(endpoint_url=self.url)
48+
49+
@patch(
50+
"codecarbon.output_methods.http.requests.post",
51+
return_value=MagicMock(status_code=201),
52+
)
53+
def test_http_output_post_success(self, mock_post):
54+
self.http_output.out(self.emissions_data, self.emissions_data)
55+
56+
mock_post.assert_called_once()
57+
self.assertEqual(mock_post.call_args[0][0], self.url)
58+
59+
@patch("codecarbon.output_methods.http.logger.warning")
60+
@patch(
61+
"codecarbon.output_methods.http.requests.post",
62+
return_value=MagicMock(status_code=418),
63+
)
64+
def test_http_output_post_unexpected_status(self, mock_post, mock_logger):
65+
self.http_output.out(self.emissions_data, self.emissions_data)
66+
67+
mock_post.assert_called_once()
68+
mock_logger.assert_called_once()
69+
70+
@patch("codecarbon.output_methods.http.logger.error")
71+
@patch(
72+
"codecarbon.output_methods.http.requests.post",
73+
side_effect=Exception("Test exception"),
74+
)
75+
def test_http_output_post_exception(self, mock_post, mock_logger):
76+
self.http_output.out(self.emissions_data, self.emissions_data)
77+
mock_post.assert_called_once()
78+
mock_logger.assert_called_once()
79+
80+
81+
class TestCodeCarbonAPIOutput(unittest.TestCase):
82+
def setUp(self):
83+
self.emissions_data = EmissionsData(
84+
timestamp="2023-01-01T00:00:00",
85+
project_name="test_project",
86+
run_id="test_run_id",
87+
experiment_id="test_experiment_id",
88+
duration=10,
89+
emissions=0.5,
90+
emissions_rate=0.05,
91+
cpu_power=20,
92+
gpu_power=30,
93+
ram_power=5,
94+
cpu_energy=200,
95+
gpu_energy=300,
96+
ram_energy=50,
97+
energy_consumed=550,
98+
water_consumed=0.1,
99+
country_name="Testland",
100+
country_iso_code="TS",
101+
region="Test Region",
102+
cloud_provider="Test Cloud",
103+
cloud_region="test-cloud-1",
104+
os="TestOS",
105+
python_version="3.8",
106+
codecarbon_version="2.0",
107+
cpu_count=4,
108+
cpu_model="Test CPU",
109+
gpu_count=1,
110+
gpu_model="Test GPU",
111+
longitude=0,
112+
latitude=0,
113+
ram_total_size=16,
114+
tracking_mode="machine",
115+
on_cloud="true",
116+
pue=1.5,
117+
wue=0.5,
118+
)
119+
self.url = "http://test.com/emissions"
120+
self.experiment_id = (
121+
None # Set to None so that ApiClient won't attempt a run on initialisation
122+
)
123+
self.api_key = "test_key"
124+
125+
self.add_emission_patcher = patch(
126+
"codecarbon.output_methods.http.ApiClient.add_emission"
127+
)
128+
self.mock_add_emission = self.add_emission_patcher.start()
129+
self.addCleanup(self.add_emission_patcher.stop)
130+
131+
def test_codecarbon_api_output_initialization(self):
132+
CodeCarbonAPIOutput(
133+
endpoint_url=self.url,
134+
experiment_id=self.experiment_id,
135+
api_key=self.api_key,
136+
conf=None,
137+
)
138+
139+
def test_codecarbon_api_live_out(self):
140+
api_output = CodeCarbonAPIOutput(
141+
endpoint_url=self.url,
142+
experiment_id=self.experiment_id,
143+
api_key=self.api_key,
144+
conf=None,
145+
)
146+
147+
api_output.live_out(None, self.emissions_data)
148+
self.mock_add_emission.assert_called_once()
149+
150+
@patch("codecarbon.output_methods.http.logger.error")
151+
def test_codecarbon_live_out_api_call_failure(self, mock_logger):
152+
self.mock_add_emission.side_effect = Exception("Test exception")
153+
api_output = CodeCarbonAPIOutput(
154+
endpoint_url=self.url,
155+
experiment_id=self.experiment_id,
156+
api_key=self.api_key,
157+
conf=None,
158+
)
159+
api_output.live_out(None, self.emissions_data)
160+
mock_logger.assert_called_once()
161+
162+
def test_codecarbon_api_out(self):
163+
api_output = CodeCarbonAPIOutput(
164+
endpoint_url=self.url,
165+
experiment_id=self.experiment_id,
166+
api_key=self.api_key,
167+
conf=None,
168+
)
169+
170+
api_output.out(None, self.emissions_data)
171+
self.mock_add_emission.assert_called_once()
172+
173+
@patch("codecarbon.output_methods.http.logger.error")
174+
def test_codecarbon_out_api_call_failure(self, mock_logger):
175+
self.mock_add_emission.side_effect = Exception("Test exception")
176+
api_output = CodeCarbonAPIOutput(
177+
endpoint_url=self.url,
178+
experiment_id=self.experiment_id,
179+
api_key=self.api_key,
180+
conf=None,
181+
)
182+
api_output.out(None, self.emissions_data)
183+
mock_logger.assert_called_once()
Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import sys
2+
import unittest
3+
from unittest.mock import patch
4+
5+
import logfire
6+
7+
from codecarbon.output_methods.emissions_data import EmissionsData
8+
from codecarbon.output_methods.metrics.logfire import LogfireOutput
9+
10+
11+
class TestLogfireOutput(unittest.TestCase):
12+
def setUp(self):
13+
self.emissions_data = EmissionsData(
14+
timestamp="2023-01-01T00:00:00",
15+
project_name="test_project",
16+
run_id="test_run_id",
17+
experiment_id="test_experiment_id",
18+
duration=10,
19+
emissions=0.5,
20+
emissions_rate=0.05,
21+
cpu_power=20,
22+
gpu_power=30,
23+
ram_power=5,
24+
cpu_energy=200,
25+
gpu_energy=300,
26+
ram_energy=50,
27+
energy_consumed=550,
28+
water_consumed=0.1,
29+
country_name="Testland",
30+
country_iso_code="TS",
31+
region="Test Region",
32+
cloud_provider="Test Cloud",
33+
cloud_region="test-cloud-1",
34+
os="TestOS",
35+
python_version="3.8",
36+
codecarbon_version="2.0",
37+
cpu_count=4,
38+
cpu_model="Test CPU",
39+
gpu_count=1,
40+
gpu_model="Test GPU",
41+
longitude=0,
42+
latitude=0,
43+
ram_total_size=16,
44+
tracking_mode="machine",
45+
on_cloud="true",
46+
pue=1.5,
47+
wue=0.5,
48+
)
49+
50+
self.configure_patcher = patch("logfire.configure")
51+
self.mock_configure = self.configure_patcher.start()
52+
self.mock_configure.return_value = logfire.configure(send_to_logfire=False)
53+
self.addCleanup(self.configure_patcher.stop)
54+
55+
@patch.dict(sys.modules, {"logfire": None})
56+
def test_logfire_import_error(self):
57+
with self.assertRaises(ImportError):
58+
LogfireOutput()
59+
60+
def test_logfire_initialization(self):
61+
LogfireOutput()
62+
63+
def _check_output_is_not_empty(self, logfire_output):
64+
self.assertIsNotNone(logfire_output.duration)
65+
self.assertIsNotNone(logfire_output.emissions)
66+
self.assertIsNotNone(logfire_output.emissions_rate)
67+
self.assertIsNotNone(logfire_output.cpu_power)
68+
self.assertIsNotNone(logfire_output.gpu_power)
69+
self.assertIsNotNone(logfire_output.ram_power)
70+
self.assertIsNotNone(logfire_output.cpu_energy)
71+
self.assertIsNotNone(logfire_output.gpu_energy)
72+
self.assertIsNotNone(logfire_output.ram_energy)
73+
self.assertIsNotNone(logfire_output.energy_consumed)
74+
75+
def test_logfire_out(self):
76+
logfire_output = LogfireOutput()
77+
logfire_output.out(self.emissions_data, self.emissions_data)
78+
self._check_output_is_not_empty(logfire_output)
79+
80+
def test_logfire_live_out(self):
81+
logfire_output = LogfireOutput()
82+
logfire_output.live_out(self.emissions_data, self.emissions_data)
83+
self._check_output_is_not_empty(logfire_output)

uv.lock

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)