From dcf02ef6c761da48cb195216f9db6f09cbf2f023 Mon Sep 17 00:00:00 2001 From: Shovnik Bhattacharya Date: Thu, 26 Nov 2020 19:05:08 -0500 Subject: [PATCH] Made changes suggested in Diego's review --- .../prometheus_remote_write/__init__.py | 36 +++++------ .../test_prometheus_remote_write_exporter.py | 62 +++++++++---------- 2 files changed, 44 insertions(+), 54 deletions(-) diff --git a/exporter/opentelemetry-exporter-prometheus-remote-write/src/opentelemetry/exporter/prometheus_remote_write/__init__.py b/exporter/opentelemetry-exporter-prometheus-remote-write/src/opentelemetry/exporter/prometheus_remote_write/__init__.py index d93ab85914..1763ef4abb 100644 --- a/exporter/opentelemetry-exporter-prometheus-remote-write/src/opentelemetry/exporter/prometheus_remote_write/__init__.py +++ b/exporter/opentelemetry-exporter-prometheus-remote-write/src/opentelemetry/exporter/prometheus_remote_write/__init__.py @@ -91,15 +91,15 @@ def basic_auth(self, basic_auth: Dict): if basic_auth: if "username" not in basic_auth: raise ValueError("username required in basic_auth") - if ( - "password" not in basic_auth - and "password_file" not in basic_auth - ): + if "password_file" in basic_auth: + if "password" in basic_auth: + raise ValueError( + "basic_auth cannot contain password and password_file" + ) + with open(basic_auth["password_file"]) as file: + basic_auth["password"] = file.readline().strip() + elif "password" not in basic_auth: raise ValueError("password required in basic_auth") - if "password" in basic_auth and "password_file" in basic_auth: - raise ValueError( - "basic_auth cannot contain password and password_file" - ) self._basic_auth = basic_auth @property @@ -154,10 +154,10 @@ def headers(self, headers: Dict): def export( self, export_records: Sequence[ExportRecord] ) -> MetricsExportResult: - timeseries = self.convert_to_timeseries(export_records) - message = self.build_message(timeseries) - headers = self.get_headers() - return self.send_message(message, headers) + timeseries = self._convert_to_timeseries(export_records) + message = self._build_message(timeseries) + headers = self._build_headers() + return self._send_message(message, headers) def shutdown(self) -> None: pass @@ -301,13 +301,13 @@ def add_label(label_name, label_value): timeseries.samples.append(sample) return timeseries - def build_message(self, timeseries: Sequence[TimeSeries]) -> bytes: + def _build_message(self, timeseries: Sequence[TimeSeries]) -> bytes: write_request = WriteRequest() write_request.timeseries.extend(timeseries) serialized_message = write_request.SerializeToString() return snappy.compress(serialized_message) - def get_headers(self) -> Dict: + def _build_headers(self) -> Dict: headers = { "Content-Encoding": "snappy", "Content-Type": "application/x-protobuf", @@ -318,16 +318,12 @@ def get_headers(self) -> Dict: headers[header_name] = header_value return headers - def send_message( + def _send_message( self, message: bytes, headers: Dict ) -> MetricsExportResult: auth = None if self.basic_auth: - if "password" in self.basic_auth: - auth = (self.basic_auth["username"], self.basic_auth["password"]) - elif "password_file": - with open(self.basic_auth["password_file"]) as file: - auth = (self.basic_auth["username"], file.readline()) + auth = (self.basic_auth["username"], self.basic_auth["password"]) cert = None verify = True diff --git a/exporter/opentelemetry-exporter-prometheus-remote-write/tests/test_prometheus_remote_write_exporter.py b/exporter/opentelemetry-exporter-prometheus-remote-write/tests/test_prometheus_remote_write_exporter.py index 809d5494eb..a0ac387c8d 100644 --- a/exporter/opentelemetry-exporter-prometheus-remote-write/tests/test_prometheus_remote_write_exporter.py +++ b/exporter/opentelemetry-exporter-prometheus-remote-write/tests/test_prometheus_remote_write_exporter.py @@ -358,65 +358,59 @@ def create_label(name, value): self.assertEqual(timeseries, expected_timeseries) -class ResponseStub: - def __init__(self, status_code): - self.status_code = status_code - self.reason = "dummy_reason" - self.content = "dummy_content" - - class TestExport(unittest.TestCase): # Initializes test data that is reused across tests def setUp(self): - self._exporter = PrometheusRemoteWriteMetricsExporter( + self.exporter = PrometheusRemoteWriteMetricsExporter( endpoint="/prom/test_endpoint" ) # Ensures export is successful with valid export_records and config - @mock.patch("requests.post", return_value=ResponseStub(200)) + @patch("requests.post") def test_export(self, mock_post): + mock_post.return_value.configure_mock(**{"status_code": 200}) test_metric = Counter("testname", "testdesc", "testunit", int, None) - labels = {"environment": "testing"} + labels = get_dict_as_key({"environment": "testing"}) record = ExportRecord( - test_metric, labels, SumAggregator(), Resource({}), + test_metric, labels, SumAggregator(), Resource({}) ) - result = self._exporter.export([record]) + result = self.exporter.export([record]) self.assertIs(result, MetricsExportResult.SUCCESS) self.assertEqual(mock_post.call_count, 1) - @mock.patch("requests.post", return_value=ResponseStub(200)) + @patch("requests.post") def test_valid_send_message(self, mock_post): - result = self._exporter.send_message(bytes(), {}) + mock_post.return_value.configure_mock(**{"status_code": 200}) + result = self.exporter._send_message(bytes(), {}) self.assertEqual(mock_post.call_count, 1) self.assertEqual(result, MetricsExportResult.SUCCESS) - @mock.patch("requests.post", return_value=ResponseStub(404)) + @patch("requests.post") def test_invalid_send_message(self, mock_post): - result = self._exporter.send_message(bytes(), {}) + mock_post.return_value.configure_mock( + **{ + "status_code": 404, + "reason": "test_reason", + "content": "test_content", + } + ) + result = self.exporter._send_message(bytes(), {}) self.assertEqual(mock_post.call_count, 1) self.assertEqual(result, MetricsExportResult.FAILURE) # Verifies that build_message calls snappy.compress and returns SerializedString - @mock.patch("snappy.compress", return_value=bytes()) + @patch("snappy.compress", return_value=bytes()) def test_build_message(self, mock_compress): - test_timeseries = [ - TimeSeries(), - TimeSeries(), - ] - message = self._exporter.build_message(test_timeseries) + message = self.exporter._build_message([TimeSeries()]) self.assertEqual(mock_compress.call_count, 1) self.assertIsInstance(message, bytes) # Ensure correct headers are added when valid config is provided - def test_get_headers(self): - self._exporter.headers = {"Custom Header": "test_header"} - - headers = self._exporter.get_headers() - self.assertEqual(headers.get("Content-Encoding", ""), "snappy") - self.assertEqual( - headers.get("Content-Type", ""), "application/x-protobuf" - ) - self.assertEqual( - headers.get("X-Prometheus-Remote-Write-Version", ""), "0.1.0" - ) - self.assertEqual(headers.get("Custom Header", ""), "test_header") + def test_build_headers(self): + self.exporter.headers = {"Custom Header": "test_header"} + + headers = self.exporter._build_headers() + self.assertEqual(headers["Content-Encoding"], "snappy") + self.assertEqual(headers["Content-Type"], "application/x-protobuf") + self.assertEqual(headers["X-Prometheus-Remote-Write-Version"], "0.1.0") + self.assertEqual(headers["Custom Header"], "test_header")