def assert_metrics_indexer_worker(self,
producer,
metrics_payload=payload,
flush_return_value=0,
with_exception=False):
producer.produce = MagicMock()
producer.flush = MagicMock(return_value=flush_return_value)
metrics_worker = MetricsIndexerWorker(producer=producer)
mock_message = Mock()
mock_message.value = MagicMock(
return_value=json.dumps(metrics_payload))
parsed = metrics_worker.process_message(mock_message)
assert parsed["tags"] == {
PGStringIndexer().resolve(string=k):
PGStringIndexer().resolve(string=str(v))
for k, v in payload["tags"].items()
}
assert parsed["metric_id"] == PGStringIndexer().resolve(
string=payload["name"])
if with_exception:
with pytest.raises(Exception,
match="didn't get all the callbacks: 1 left"):
metrics_worker.flush_batch([parsed])
else:
metrics_worker.flush_batch([parsed])
producer.produce.assert_called_with(
topic="snuba-metrics",
key=None,
value=json.dumps(parsed).encode(),
on_delivery=metrics_worker.callback,
)
def test_metrics_consumer(self):
ingest_producer = self._get_producer(self.ingest_topic)
message = json.dumps(payload).encode()
# produce message to the dummy ingest-metrics topic
ingest_producer.produce(self.ingest_topic, message)
assert ingest_producer.flush() == 0
options = {
"max_batch_size": 1,
"max_batch_time": 5000,
"group_id": "test-metrics-indexer-consumer",
"auto_offset_reset": "earliest",
}
batching_consumer = get_metrics_consumer(topic=self.ingest_topic,
**options)
# couldn't use _run_once() here because .poll() is called
# with a 1 second timeout which seems to not be enough.
msg = batching_consumer.consumer.poll(5)
assert msg
# _handle_message calls worker's process_message
# and then we flush() to make sure we call flush_batch
batching_consumer._handle_message(msg)
batching_consumer._flush()
# make sure we produced the message during flush_batch
snuba_producer = batching_consumer.worker._MetricsIndexerWorker__producer
assert snuba_producer.flush() == 0
# in order to test that the message we produced to the dummy
# snuba-metrics topic was the message we expected, we make a
# dummy consumer to subscribe to the topic
snuba_metrics_consumer = Consumer({
"bootstrap.servers": "localhost:9092",
"group.id": "test-snuba-metrics-consumer",
"default.topic.config": {
"auto.offset.reset": "earliest"
},
})
snuba_metrics_consumer.subscribe([self.snuba_topic])
# once we have the message, we don't need the consumer anymore
translated_msg = snuba_metrics_consumer.poll(5)
snuba_metrics_consumer.close()
assert translated_msg
# finally test the payload of the translated message
parsed = json.loads(translated_msg.value(), use_rapid_json=True)
assert parsed["tags"] == {
str(PGStringIndexer().resolve(string=k)):
PGStringIndexer().resolve(string=str(v))
for k, v in payload["tags"].items()
}
assert parsed["metric_id"] == PGStringIndexer().resolve(
string=payload["name"])
def translate_payload():
parsed = payload.copy()
parsed["tags"] = {
PGStringIndexer().resolve(string=k):
PGStringIndexer().resolve(string=str(v))
for k, v in payload["tags"].items()
}
parsed["metric_id"] = PGStringIndexer().resolve(string=payload["name"])
# hard-coded retention days added in by the consumer
parsed["retention_days"] = 90
return parsed
def test_indexer(self):
results = PGStringIndexer().bulk_record(strings=["hello", "hey", "hi"])
assert list(results.values()) == [1, 2, 3]
# test resolve and reverse_resolve
obj = MetricsKeyIndexer.objects.get(string="hello")
assert PGStringIndexer().resolve("hello") == obj.id
assert PGStringIndexer().reverse_resolve(obj.id) == obj.string
# test record on a string that already exists
PGStringIndexer().record("hello")
assert PGStringIndexer().resolve("hello") == obj.id
# test invalid values
assert PGStringIndexer().resolve("beep") is None
assert PGStringIndexer().reverse_resolve(1234) is None
def test_indexer(self):
org_id = self.organization.id
org_strings = {org_id: {"hello", "hey", "hi"}}
results = PGStringIndexer().bulk_record(org_strings=org_strings)
obj_ids = list(
MetricsKeyIndexer.objects.filter(
string__in=["hello", "hey", "hi"]).values_list("id",
flat=True))
assert list(results.values()) == obj_ids
# test resolve and reverse_resolve
obj = MetricsKeyIndexer.objects.get(string="hello")
assert PGStringIndexer().resolve("hello") == obj.id
assert PGStringIndexer().reverse_resolve(obj.id) == obj.string
# test record on a string that already exists
PGStringIndexer().record(org_id, "hello")
assert PGStringIndexer().resolve("hello") == obj.id
# test invalid values
assert PGStringIndexer().resolve("beep") is None
assert PGStringIndexer().reverse_resolve(1234) is None
def setUp(self) -> None:
self.indexer = PGStringIndexer()