Adding exporter methods

exporter/opentelemetry-exporter-prometheus-remote-write/src/opentelemetry/exporter/prometheus_remote_write/__init__.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 import re
 from typing import Dict, Sequence
 
@@ -35,6 +36,8 @@
     ValueObserverAggregator,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class PrometheusRemoteWriteMetricsExporter(MetricsExporter):
     """
@@ -265,12 +268,50 @@ def create_label(self, name: str, value: str) -> Label:
         return label
 
     def build_message(self, timeseries: Sequence[TimeSeries]) -> bytes:
-        raise NotImplementedError()
+        print(timeseries)
+        write_request = WriteRequest()
+        write_request.timeseries.extend(timeseries)
+        serialized_message = write_request.SerializeToString()
+        return snappy.compress(serialized_message)
 
     def get_headers(self) -> Dict:
-        raise NotImplementedError()
+        headers = {
+            "Content-Encoding": "snappy",
+            "Content-Type": "application/x-protobuf",
+            "X-Prometheus-Remote-Write-Version": "0.1.0",
+        }
+        if hasattr(self, "headers"):
+            for header_name, header_value in self.headers.items():
+                headers[header_name] = header_value
+
+        if "Authorization" not in headers:
+            if hasattr(self, "bearer_token"):
+                headers["Authorization"] = "Bearer " + self.bearer_token
+            elif hasattr(self, "bearer_token_file"):
+                with open(self.bearer_token_file) as file:
+                    headers["Authorization"] = "Bearer " + file.readline()
+        return headers
 
     def send_message(
         self, message: bytes, headers: Dict
     ) -> MetricsExportResult:
-        raise NotImplementedError()
+        auth = None
+        if hasattr(self, "basic_auth"):
+            basic_auth = self.basic_auth
+            if "password" in basic_auth:
+                auth = (basic_auth.username, basic_auth.password)
+            else:
+                with open(basic_auth.password_file) as file:
+                    auth = (basic_auth.username, file.readline())
+        response = requests.post(
+            self.endpoint, data=message, headers=headers, auth=auth
+        )
+        if response.status_code != 200:
+            logger.warning(
+                "POST request failed with status %s with reason: %s and content: %s",
+                str(response.status_code),
+                response.reason,
+                str(response.content),
+            )
+            return MetricsExportResult.FAILURE
+        return MetricsExportResult.SUCCESS

exporter/opentelemetry-exporter-prometheus-remote-write/tests/test_prometheus_remote_write_exporter.py

@@ -288,25 +288,74 @@ def test_create_timeseries(self):
         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):
-        pass
+        self._exporter = PrometheusRemoteWriteMetricsExporter(
+            endpoint="/prom/test_endpoint"
+        )
 
     # Ensures export is successful with valid export_records and config
     def test_export(self):
-        pass
+        test_metric = Counter("testname", "testdesc", "testunit", int, None)
+        record = ExportRecord(
+            test_metric,
+            None,
+            SumAggregator(),
+            Resource({}),
+        )
+        self._exporter.convert_to_timeseries = mock.Mock(return_value=[])
+        self._exporter.build_message = mock.Mock(return_value=bytes())
+        self._exporter.get_headers = mock.Mock(return_value={})
+        self._exporter.send_message = mock.Mock(
+            return_value=MetricsExportResult.SUCCESS
+        )
+        result = self._exporter.export([record])
+        self.assertIs(result, MetricsExportResult.SUCCESS)
 
-    def test_valid_send_message(self):
-        pass
+    @mock.patch("requests.post", return_value=ResponseStub(200))
+    def test_valid_send_message(self, mock_post):
+        result = self._exporter.send_message(bytes(), {})
+        self.assertEqual(mock_post.call_count, 1)
+        self.assertEqual(result, MetricsExportResult.SUCCESS)
 
-    def test_invalid_send_message(self):
-        pass
+    @mock.patch("requests.post", return_value=ResponseStub(404))
+    def test_invalid_send_message(self, mock_post):
+        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
-    def test_build_message(self):
-        pass
+    @mock.patch("snappy.compress", return_value=bytes())
+    def test_build_message(self, mock_compress):
+        test_timeseries = [
+            TimeSeries(),
+            TimeSeries(),
+        ]
+        message = self._exporter.build_message(test_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):
-        pass
+        self._exporter.headers = {"Custom Header": "test_header"}
+        self._exporter.headers = {"Custom Header": "test_header"}
+        self._exporter.bearer_token = "test_token"
+
+        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("Authorization", ""), "Bearer test_token")
+        self.assertEqual(headers.get("Custom Header", ""), "test_header")