def test_normal(self):
cluster_name = settings.KAFKA_TOPICS[self.topic]["cluster"]
conf = {
"bootstrap.servers":
settings.KAFKA_CLUSTERS[cluster_name]["bootstrap.servers"],
"session.timeout.ms":
6000,
}
producer = Producer(conf)
producer.produce(self.topic, json.dumps(self.valid_wrapper))
producer.flush()
mock_callback = Mock()
mock_callback.side_effect = KeyboardInterrupt()
register_subscriber(self.registration_key)(mock_callback)
sub = QuerySubscription.objects.create(
project=self.project,
type=self.registration_key,
subscription_id=self.subscription_id,
dataset="something",
query="hello",
aggregation=0,
time_window=1,
resolution=1,
)
consumer = QuerySubscriptionConsumer("hi",
topic=self.topic,
commit_batch_size=1)
consumer.run()
mock_callback.assert_called_once_with(self.valid_payload, sub)
def test_normal(self):
cluster_name = settings.KAFKA_TOPICS[self.topic]["cluster"]
conf = {
"bootstrap.servers":
settings.KAFKA_CLUSTERS[cluster_name]["bootstrap.servers"],
"session.timeout.ms":
6000,
}
producer = Producer(conf)
producer.produce(self.topic, json.dumps(self.valid_wrapper))
producer.flush()
mock_callback = Mock()
mock_callback.side_effect = KeyboardInterrupt()
register_subscriber(self.registration_key)(mock_callback)
sub = self.create_subscription()
consumer = QuerySubscriptionConsumer("hi",
topic=self.topic,
commit_batch_size=1)
consumer.run()
payload = self.valid_payload
payload["values"] = payload["result"]
payload["timestamp"] = parse_date(
payload["timestamp"]).replace(tzinfo=pytz.utc)
mock_callback.assert_called_once_with(payload, sub)
def test_shutdown(self):
self.producer.produce(self.topic, json.dumps(self.valid_wrapper))
valid_wrapper_2 = deepcopy(self.valid_wrapper)
valid_wrapper_2["payload"]["values"]["hello"] = 25
valid_wrapper_3 = deepcopy(valid_wrapper_2)
valid_wrapper_3["payload"]["values"]["hello"] = 5000
self.producer.produce(self.topic, json.dumps(valid_wrapper_2))
self.producer.flush()
counts = [0]
def mock_callback(*args, **kwargs):
counts[0] += 1
if counts[0] > 1:
raise KeyboardInterrupt()
mock = Mock()
mock.side_effect = mock_callback
register_subscriber(self.registration_key)(mock)
sub = QuerySubscription.objects.create(
project=self.project,
type=self.registration_key,
subscription_id=self.subscription_id,
dataset="something",
query="hello",
aggregation=0,
time_window=1,
resolution=1,
)
consumer = QuerySubscriptionConsumer("hi",
topic=self.topic,
commit_batch_size=100)
consumer.run()
valid_payload = self.valid_payload
valid_payload["timestamp"] = parse_date(
valid_payload["timestamp"]).replace(tzinfo=pytz.utc)
valid_wrapper_2["payload"]["timestamp"] = parse_date(
valid_wrapper_2["payload"]["timestamp"]).replace(tzinfo=pytz.utc)
mock.assert_has_calls(
[call(valid_payload, sub),
call(valid_wrapper_2["payload"], sub)])
# Offset should be committed for the first message, so second run should process
# the second message again
self.producer.produce(self.topic, json.dumps(valid_wrapper_3))
self.producer.flush()
mock.reset_mock()
counts[0] = 0
consumer.run()
valid_wrapper_3["payload"]["timestamp"] = parse_date(
valid_wrapper_3["payload"]["timestamp"]).replace(tzinfo=pytz.utc)
mock.assert_has_calls([
call(valid_wrapper_2["payload"], sub),
call(valid_wrapper_3["payload"], sub)
])
