159 lines
7.1 KiB
Python
159 lines
7.1 KiB
Python
"""Worker lifecycle and retry behavior without broker or database processes."""
|
|
import json
|
|
import signal
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, call, patch
|
|
|
|
from app.api.producer import MockMessageProducer
|
|
from app.worker.clickhouse import ClickHouseService
|
|
from app.worker.consumer import ProcessingWorker
|
|
from app.worker.metrics import WorkerMetrics
|
|
from app.worker.normalizer import normalize_cdr_payload
|
|
from tests.test_delivery_regressions import Message, payload
|
|
|
|
|
|
class WorkerLifecycleTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.consumer = Mock()
|
|
self.consumer.commit.return_value = []
|
|
self.consumer.assignment.return_value = []
|
|
self.producer = MockMessageProducer()
|
|
self.store = Mock()
|
|
self.server = Mock()
|
|
self.metrics = WorkerMetrics()
|
|
metric_patch = patch("app.worker.consumer.worker_metrics", self.metrics)
|
|
metric_patch.start()
|
|
self.addCleanup(metric_patch.stop)
|
|
self.worker = ProcessingWorker(consumer=self.consumer, dlq_producer=self.producer,
|
|
clickhouse=self.store, batch_size=2,
|
|
flush_interval=1, internal_server=self.server)
|
|
|
|
def test_count_trigger_flushes_exact_batch_before_shutdown(self):
|
|
self.worker.process_single_message(Message(payload(), offset=0))
|
|
self.store.insert_batch.assert_not_called()
|
|
self.worker.process_single_message(Message(payload(raw_id="second"), offset=1))
|
|
self.assertEqual(len(self.store.insert_batch.call_args.args[0]), 2)
|
|
self.assertEqual(self.worker.buffer, [])
|
|
self.assertEqual(self.consumer.commit.call_args.kwargs["offsets"][0].offset, 2)
|
|
|
|
def test_idle_poll_flushes_on_timer_without_waiting_for_shutdown(self):
|
|
now = [0.0]
|
|
self.worker.last_flush_time = 0.0
|
|
polled = [0]
|
|
|
|
def poll(timeout):
|
|
polled[0] += 1
|
|
if polled[0] == 1:
|
|
return Message(payload())
|
|
if polled[0] == 2:
|
|
now[0] = 1.1
|
|
return None
|
|
self.store.insert_batch.assert_called_once()
|
|
self.assertEqual(self.worker.buffer, [])
|
|
self.worker._handle_signal(signal.SIGTERM, None)
|
|
return None
|
|
|
|
self.consumer.poll.side_effect = poll
|
|
with patch("app.worker.consumer.time.monotonic", side_effect=lambda: now[0]), \
|
|
patch("app.worker.consumer.signal.signal"):
|
|
self.worker.start()
|
|
self.store.insert_batch.assert_called_once()
|
|
self.consumer.close.assert_called_once()
|
|
self.server.stop.assert_called_once()
|
|
|
|
def test_graceful_shutdown_flushes_partial_batch(self):
|
|
self.worker.process_single_message(Message(payload(), offset=3))
|
|
self.worker.stop()
|
|
self.store.insert_batch.assert_called_once()
|
|
self.assertEqual(self.consumer.commit.call_args.kwargs["offsets"][0].offset, 4)
|
|
self.consumer.close.assert_called_once()
|
|
self.store.close.assert_called_once()
|
|
self.server.stop.assert_called_once()
|
|
|
|
def test_failed_commit_cleanup_does_not_retry_write_or_commit(self):
|
|
self.consumer.commit.side_effect = ConnectionError("commit failed")
|
|
self.worker.process_single_message(Message(payload()))
|
|
with self.assertRaises(ConnectionError):
|
|
self.worker.flush_buffer()
|
|
self.worker.stop()
|
|
self.assertTrue(self.worker.failed)
|
|
self.store.insert_batch.assert_called_once()
|
|
self.consumer.commit.assert_called_once()
|
|
self.assertEqual(self.producer.messages, [])
|
|
self.consumer.close.assert_called_once()
|
|
self.store.close.assert_called_once()
|
|
self.server.stop.assert_called_once()
|
|
|
|
def test_lag_uses_committed_offsets_and_log_start_for_uncommitted_partition(self):
|
|
partitions = [SimpleNamespace(partition=0, offset=7),
|
|
SimpleNamespace(partition=1, offset=-1001)]
|
|
self.consumer.assignment.return_value = partitions
|
|
self.consumer.committed.return_value = partitions
|
|
self.consumer.get_watermark_offsets.side_effect = [(0, 12), (5, 20)]
|
|
self.worker.update_consumer_lag()
|
|
output = self.metrics.get_prometheus_metrics()
|
|
self.assertIn('cdr_kafka_consumer_lag{partition="0"} 5', output)
|
|
self.assertIn('cdr_kafka_consumer_lag{partition="1"} 15', output)
|
|
|
|
def test_replay_after_write_before_commit_preserves_logical_identity(self):
|
|
# Models FINAL replacement by identity; real sorting-key verification is integration coverage.
|
|
logical_rows = {}
|
|
writes = []
|
|
|
|
def insert(records):
|
|
writes.extend(records)
|
|
logical_rows.update((record.cdr_id, record) for record in records)
|
|
|
|
self.store.insert_batch.side_effect = insert
|
|
self.consumer.commit.side_effect = ConnectionError("crash before commit")
|
|
source = Message(payload(), offset=8)
|
|
self.worker.process_single_message(source)
|
|
with self.assertRaises(ConnectionError):
|
|
self.worker.flush_buffer()
|
|
replay_consumer = Mock()
|
|
replay_consumer.commit.return_value = []
|
|
restarted = ProcessingWorker(consumer=replay_consumer, dlq_producer=self.producer,
|
|
clickhouse=self.store, batch_size=1, internal_server=Mock())
|
|
restarted.process_single_message(source)
|
|
self.assertEqual(len(writes), 2)
|
|
self.assertEqual(len(logical_rows), 1)
|
|
self.assertEqual(replay_consumer.commit.call_args.kwargs["offsets"][0].offset, 9)
|
|
self.assertEqual(self.producer.messages, [])
|
|
|
|
|
|
class ClickHouseRetryTests(unittest.TestCase):
|
|
def make_service(self):
|
|
service = ClickHouseService.__new__(ClickHouseService)
|
|
service._client = Mock()
|
|
service.table = "cdr_records"
|
|
service.max_attempts = 3
|
|
service.initial_delay_ms = 100
|
|
service.backoff_factor = 2
|
|
record, reason, _, _ = normalize_cdr_payload(json.dumps(payload()).encode())
|
|
self.assertIsNone(reason)
|
|
return service, record
|
|
|
|
def test_transient_failures_retry_with_exponential_delays(self):
|
|
service, record = self.make_service()
|
|
service._client.insert.side_effect = [TimeoutError("first"), TimeoutError("second"), None]
|
|
with patch("app.worker.clickhouse.time.sleep") as sleep:
|
|
self.assertTrue(service.insert_batch([record]))
|
|
self.assertEqual(service._client.insert.call_count, 3)
|
|
self.assertEqual(sleep.call_args_list, [call(0.1), call(0.2)])
|
|
|
|
def test_exhausted_retries_raise_last_failure_without_final_sleep(self):
|
|
service, record = self.make_service()
|
|
failure = TimeoutError("persistent failure")
|
|
service._client.insert.side_effect = failure
|
|
with patch("app.worker.clickhouse.time.sleep") as sleep:
|
|
with self.assertRaises(TimeoutError) as caught:
|
|
service.insert_batch([record])
|
|
self.assertIs(caught.exception, failure)
|
|
self.assertEqual(service._client.insert.call_count, 3)
|
|
self.assertEqual(sleep.call_args_list, [call(0.1), call(0.2)])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|