From 3ab4d8961a7421ef5c3cb2bbf9f94b1c430c3981 Mon Sep 17 00:00:00 2001 From: sharonsyh Date: Sat, 30 Nov 2024 19:47:38 -0500 Subject: [PATCH] Patch to http.client.HTTPConnection --- tests/test_metric.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index dc7e8683..be1eaf42 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -66,7 +66,7 @@ def mock_gauge(): yield gauge -@patch("prometheus_client.exposition.push_to_gateway") +@patch("prometheus_client.exposition.push_to_gateway", autospec=True) def test_energy_histogram( mock_push_to_gateway: MagicMock, mock_get_cpus: MagicMock, @@ -79,7 +79,7 @@ def test_energy_histogram( and that the correct energy values are recorded. Args: - mock_push_to_gateway (MagicMock): Mocked `push_to_gateway` function for Prometheus. + mock_push_to_gateway (MagicMock): Mocked `push_to_gateway` mock_get_cpus (MagicMock): Mocked `get_cpus` fixture. mock_zeus_monitor (MagicMock): Mocked ZeusMonitor fixture. mock_histogram (MagicMock): Mocked Prometheus Histogram fixture. @@ -111,11 +111,13 @@ def test_energy_histogram( dram_histogram.observe = MagicMock() histogram_metric.begin_window("test_window") - - with patch("prometheus_client.exposition.push_to_gateway") as mock_push: - mock_push.return_value = None + with patch("http.client.HTTPConnection", autospec=True) as mock_http: + mock_http_instance = mock_http.return_value + mock_http_instance.getresponse.return_value.code = 200 + mock_http_instance.getresponse.return_value.msg = "OK" + mock_http_instance.getresponse.return_value.info = lambda: {} + mock_http_instance.sock = MagicMock() histogram_metric.end_window("test_window") - # Assert GPU histograms were observed if mock_zeus_monitor.return_value.end_window.return_value.gpu_energy: for (