amiirkhl/telecom-cdr-pipeline-ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(worker): commit confirmed partition offsets and drain safely

Browse Source
This commit is contained in:
Amirhossein Khalili
2026-09-19 21:30:33 +03:30
parent 33593609e3
commit 3b59d76f67
5 changed files with 483 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

193
tests/test_delivery_regressions.py Normal file
View File

@@ -0,0 +1,193 @@
"""Durability regressions using broker acknowledgement and partition-aware doubles."""
import json
import time
import unittest
from concurrent.futures import Future
from types import SimpleNamespace
from unittest.mock import Mock, patch
from app.worker.consumer import ProcessingWorker
from app.worker.metrics import WorkerMetrics
from app.worker.normalizer import normalize_cdr_payload
def payload(**changes):
data = dict(raw_id="regression-1", device_id="310410998877665",
caller_num="0019875551234", svc="4G_DATA", rat=6,
up_bytes=100, down_bytes=200, lac="4021", ci="8812",
epoch_start=1726070000, epoch_end=1726070300, src_ip="10.0.0.1")
data.update(changes)
return data
class Message:
def __init__(self, data, offset=0, partition=0):
self.data = data if isinstance(data, bytes) else json.dumps(data).encode()
self.position = offset
self.partition_id = partition
def value(self): return self.data
def topic(self): return "cdr.raw"
def partition(self): return self.partition_id
def offset(self): return self.position
def timestamp(self): return (1, int((time.time() - 2) * 1000))
def error(self): return None
class Consumer:
def __init__(self, events, error=None, partition_error=False):
self.events = events
self.error = error
self.partition_error = partition_error
self.commits = []
def commit(self, offsets=None, message=None, asynchronous=False):
if self.error:
raise self.error
if offsets is not None:
positions = {(item.topic, item.partition): item.offset for item in offsets}
elif message is not None:
positions = {(message.topic(), message.partition()): message.offset() + 1}
else:
positions = "implicit-current-position"
self.events.append(("commit", positions))
self.commits.append(positions)
return [SimpleNamespace(error=RuntimeError("partition commit failed"))] if self.partition_error else []
def close(self): pass
class Producer:
def __init__(self, events, failures=()):
self.events = events
self.failures = set(failures)
self.messages = []
def produce(self, topic, key, value):
self.messages.append(json.loads(value))
result = Future()
if len(self.messages) in self.failures:
result.set_exception(ConnectionError("broker rejected DLQ record"))
else:
self.events.append(("dlq_ack", key))
result.set_result(None)
return result
def close(self): pass
def flush(self, timeout=5): return 0
class DeliverySafetyTests(unittest.TestCase):
def setUp(self):
self.events = []
self.consumer = Consumer(self.events)
self.producer = Producer(self.events)
self.ch = Mock()
self.ch.insert_batch.side_effect = lambda records: self.events.append(("insert", len(records))) or True
self.metrics = WorkerMetrics()
self.patches = [patch("app.worker.consumer.worker_metrics", self.metrics),
patch("app.worker.dlq.worker_metrics", self.metrics)]
for item in self.patches:
item.start()
self.addCleanup(item.stop)
self.worker = ProcessingWorker(consumer=self.consumer, dlq_producer=self.producer,
clickhouse=self.ch, batch_size=100, internal_server=Mock())
def test_invalid_record_cannot_commit_past_buffered_valid_record(self):
self.worker.process_single_message(Message(payload(), offset=0))
self.worker.process_single_message(Message(b"{malformed", offset=1))
self.assertEqual(self.events[0], ("insert", 1))
self.assertEqual(self.consumer.commits[-1], {("cdr.raw", 0): 2})
self.assertEqual(len(self.worker.buffer), 0)
def test_batch_commits_only_its_explicit_partition_offsets(self):
self.worker.process_single_message(Message(payload(), offset=4, partition=0))
self.worker.process_single_message(Message(payload(raw_id="second"), offset=9, partition=1))
self.worker.flush_buffer()
self.assertEqual(self.consumer.commits, [{("cdr.raw", 0): 5, ("cdr.raw", 1): 10}])
def test_failed_dlq_ack_never_commits_or_counts_success(self):
self.producer.failures = {1}
with self.assertRaises(Exception):
self.worker.process_single_message(Message(b"{malformed"))
self.assertEqual(self.consumer.commits, [])
self.assertEqual(self.metrics.get_internal_stats()["pipeline"]["total_dlq"], 0)
def test_partial_batch_dlq_failure_keeps_all_offsets_uncommitted(self):
self.ch.insert_batch.side_effect = ConnectionError("ClickHouse unavailable")
self.producer.failures = {2}
self.worker.process_single_message(Message(payload(), offset=0))
self.worker.process_single_message(Message(payload(raw_id="second"), offset=1))
with self.assertRaises(Exception):
self.worker.flush_buffer()
self.assertEqual(self.consumer.commits, [])
self.assertEqual(self.metrics.get_internal_stats()["pipeline"]["total_dlq"], 1)
def test_commit_exception_propagates_without_routing_inserted_batch_to_dlq(self):
self.consumer.error = ConnectionError("commit unavailable")
self.worker.process_single_message(Message(payload()))
with self.assertRaises(Exception):
self.worker.flush_buffer()
self.ch.insert_batch.assert_called_once()
self.assertEqual(self.producer.messages, [])
def test_partition_commit_error_is_not_success(self):
self.consumer.partition_error = True
self.worker.process_single_message(Message(payload()))
with self.assertRaises(Exception):
self.worker.flush_buffer()
self.assertEqual(self.producer.messages, [])
def test_clickhouse_failure_dlq_retains_original_vendor_payload(self):
self.ch.insert_batch.side_effect = ConnectionError("ClickHouse unavailable")
original = payload(vendor_extension={"keep": True})
self.worker.process_single_message(Message(original))
self.worker.flush_buffer()
received = self.producer.messages[0]["original_payload"]
self.assertEqual(json.loads(received) if isinstance(received, str) else received, original)
def test_successful_insert_populates_kafka_latency(self):
self.worker.process_single_message(Message(payload()))
self.worker.flush_buffer()
self.assertIn("cdr_kafka_to_clickhouse_latency_seconds_count 1", self.metrics.get_prometheus_metrics())
class InputBoundaryTests(unittest.TestCase):
def test_invalid_values_are_rejected_without_raising(self):
cases = [("epoch_start", 10**100), ("epoch_end", 10**100), ("epoch_end", 2**32),
("epoch_start", -1), ("epoch_start", 1726070000.5),
("rat", 6.9), ("rat", True), ("up_bytes", -0.5),
("up_bytes", True), ("up_bytes", 2**64),
("up_bytes", 2**64 - 1), ("src_ip", "999.1.2.3"),
("src_ip", "::1"), ("raw_id", 123)]
for field, value in cases:
with self.subTest(field=field, value=value):
record, reason, failed_field, _ = normalize_cdr_payload(json.dumps(payload(**{field: value})).encode())
self.assertIsNone(record)
self.assertIsNotNone(reason)
self.assertIsNotNone(failed_field)
def test_invalid_utf8_and_nonfinite_json_are_rejected(self):
for value in (b"\xff", json.dumps(payload(up_bytes=float("nan"))).encode()):
with self.subTest(value=value):
record, reason, _, _ = normalize_cdr_payload(value)
self.assertIsNone(record)
self.assertIsNotNone(reason)
def test_clickhouse_rows_use_utc_and_native_ipv4(self):
from datetime import timedelta
from ipaddress import IPv4Address
record, reason, _, _ = normalize_cdr_payload(json.dumps(payload()).encode())
self.assertIsNone(reason)
row = record.to_row()
self.assertEqual(row[11].utcoffset(), timedelta(0))
self.assertIsInstance(row[13], IPv4Address)
def test_prometheus_labels_escape_quotes_backslashes_and_newlines(self):
metrics = WorkerMetrics()
metrics.record_processed('svc"\\\nname', "LTE_4G")
self.assertIn('service_type="svc\\"\\\\\\nname"', metrics.get_prometheus_metrics())
if __name__ == "__main__":
unittest.main()

2
tests/test_worker.py
View File

@@ -43,7 +43,7 @@ class MockKafkaConsumer:
self.committed = False
self.closed = False
def commit(self, message=None, asynchronous=False):
def commit(self, message=None, offsets=None, asynchronous=False):
self.committed = True
def close(self):

158
tests/test_worker_lifecycle.py Normal file
View File

@@ -0,0 +1,158 @@
"""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()