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

tests/test_worker_lifecycle.py

@@ -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()