Skip to content

Commit

Permalink
Made changes suggested in Diego's review
Browse files Browse the repository at this point in the history
  • Loading branch information
shovnik committed Nov 27, 2020
1 parent 55377e4 commit dcf02ef
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 54 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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")

0 comments on commit dcf02ef

Please sign in to comment.