Add more integration tests for kafka loader

Author: incrypto32

src/amp/loaders/implementations/kafka_loader.py

@@ -105,12 +105,17 @@ def _extract_message_key(self, row: Dict[str, Any]) -> Optional[bytes]:
 
         return str(key_value).encode('utf-8')
 
-    def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str) -> None:
+    def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str, connection_name: str) -> None:
         """
         Handle blockchain reorganization by sending reorg events to the same topic.
 
         Reorg events are sent as special messages with _type='reorg' so consumers
         can detect and handle invalidated block ranges.
+
+        Args:
+            invalidation_ranges: List of block ranges to invalidate
+            table_name: The Kafka topic name to send reorg events to
+            connection_name: Connection identifier (unused for Kafka, but required by base class)
         """
         if not invalidation_ranges:
             return

tests/conftest.py

@@ -210,7 +210,12 @@ def kafka_container():
     if not TESTCONTAINERS_AVAILABLE:
         pytest.skip('Testcontainers not available')
 
+    # Configure Kafka for transactions in single-broker setup
+    # These settings are required for transactional producers to work
     container = KafkaContainer(image='confluentinc/cp-kafka:7.6.0')
+    container.with_env('KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR', '1')
+    container.with_env('KAFKA_TRANSACTION_STATE_LOG_MIN_ISR', '1')
+    container.with_env('KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR', '1')
     container.start()
 
     time.sleep(10)

tests/integration/test_kafka_loader.py

@@ -1,8 +1,14 @@
+import json
+from unittest.mock import patch
+
 import pyarrow as pa
 import pytest
+from kafka import KafkaConsumer
+from kafka.errors import KafkaError
 
 try:
     from src.amp.loaders.implementations.kafka_loader import KafkaLoader
+    from src.amp.streaming.types import BatchMetadata, BlockRange, ResponseBatch
 except ImportError:
     pytest.skip('amp modules not available', allow_module_level=True)
 
@@ -42,3 +48,209 @@ def test_load_batch(self, kafka_test_config):
 
         assert result.success == True
         assert result.rows_loaded == 3
+
+    def test_message_consumption_verification(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+        topic_name = 'test_consumption_topic'
+
+        batch = pa.RecordBatch.from_pydict(
+            {
+                'id': [1, 2, 3],
+                'name': ['alice', 'bob', 'charlie'],
+                'score': [100, 200, 150],
+                'active': [True, False, True],
+            }
+        )
+
+        with loader:
+            result = loader.load_batch(batch, topic_name)
+
+        assert result.success is True
+        assert result.rows_loaded == 3
+
+        consumer = KafkaConsumer(
+            topic_name,
+            bootstrap_servers=kafka_test_config['bootstrap_servers'],
+            auto_offset_reset='earliest',
+            consumer_timeout_ms=5000,
+            value_deserializer=lambda m: json.loads(m.decode('utf-8')),
+        )
+
+        messages = list(consumer)
+        consumer.close()
+
+        assert len(messages) == 3
+
+        for i, msg in enumerate(messages):
+            assert msg.key == str(i + 1).encode('utf-8')
+            assert msg.value['_type'] == 'data'
+            assert msg.value['id'] == i + 1
+            assert msg.value['name'] in ['alice', 'bob', 'charlie']
+            assert msg.value['score'] in [100, 200, 150]
+            assert msg.value['active'] in [True, False]
+
+        msg1 = messages[0]
+        assert msg1.value['id'] == 1
+        assert msg1.value['name'] == 'alice'
+        assert msg1.value['score'] == 100
+        assert msg1.value['active'] is True
+
+        msg2 = messages[1]
+        assert msg2.value['id'] == 2
+        assert msg2.value['name'] == 'bob'
+        assert msg2.value['score'] == 200
+        assert msg2.value['active'] is False
+
+        msg3 = messages[2]
+        assert msg3.value['id'] == 3
+        assert msg3.value['name'] == 'charlie'
+        assert msg3.value['score'] == 150
+        assert msg3.value['active'] is True
+
+    def test_handle_reorg(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+        topic_name = 'test_reorg_topic'
+
+        invalidation_ranges = [
+            BlockRange(network='ethereum', start=100, end=200),
+            BlockRange(network='polygon', start=500, end=600),
+        ]
+
+        with loader:
+            loader._handle_reorg(invalidation_ranges, topic_name, 'test_connection')
+
+        consumer = KafkaConsumer(
+            topic_name,
+            bootstrap_servers=kafka_test_config['bootstrap_servers'],
+            auto_offset_reset='earliest',
+            consumer_timeout_ms=5000,
+            value_deserializer=lambda m: json.loads(m.decode('utf-8')),
+        )
+
+        messages = list(consumer)
+        consumer.close()
+
+        assert len(messages) == 2
+
+        msg1 = messages[0]
+        assert msg1.key == b'reorg:ethereum'
+        assert msg1.value['_type'] == 'reorg'
+        assert msg1.value['network'] == 'ethereum'
+        assert msg1.value['start_block'] == 100
+        assert msg1.value['end_block'] == 200
+
+        msg2 = messages[1]
+        assert msg2.key == b'reorg:polygon'
+        assert msg2.value['_type'] == 'reorg'
+        assert msg2.value['network'] == 'polygon'
+        assert msg2.value['start_block'] == 500
+        assert msg2.value['end_block'] == 600
+
+    def test_streaming_with_reorg(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+        topic_name = 'test_streaming_topic'
+
+        data1 = pa.RecordBatch.from_pydict({'id': [1, 2], 'value': [100, 200]})
+        data2 = pa.RecordBatch.from_pydict({'id': [3, 4], 'value': [300, 400]})
+        data3 = pa.RecordBatch.from_pydict({'id': [5, 6], 'value': [500, 600]})
+
+        response1 = ResponseBatch.data_batch(
+            data=data1,
+            metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=100, end=110, hash='0xabc123')]),
+        )
+
+        response2 = ResponseBatch.data_batch(
+            data=data2,
+            metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=110, end=120, hash='0xdef456')]),
+        )
+
+        reorg_response = ResponseBatch.reorg_batch(
+            invalidation_ranges=[BlockRange(network='ethereum', start=110, end=200)]
+        )
+
+        response3 = ResponseBatch.data_batch(
+            data=data3,
+            metadata=BatchMetadata(ranges=[BlockRange(network='ethereum', start=110, end=120, hash='0xnew123')]),
+        )
+
+        stream = [response1, response2, reorg_response, response3]
+
+        with loader:
+            results = list(loader.load_stream_continuous(iter(stream), topic_name))
+
+        assert len(results) == 4
+        assert results[0].success
+        assert results[0].rows_loaded == 2
+        assert results[1].success
+        assert results[1].rows_loaded == 2
+        assert results[2].success
+        assert results[2].is_reorg
+        assert results[3].success
+        assert results[3].rows_loaded == 2
+
+        consumer = KafkaConsumer(
+            topic_name,
+            bootstrap_servers=kafka_test_config['bootstrap_servers'],
+            auto_offset_reset='earliest',
+            consumer_timeout_ms=5000,
+            value_deserializer=lambda m: json.loads(m.decode('utf-8')),
+        )
+
+        messages = list(consumer)
+        consumer.close()
+
+        assert len(messages) == 7
+
+        data_messages = [msg for msg in messages if msg.value.get('_type') == 'data']
+        reorg_messages = [msg for msg in messages if msg.value.get('_type') == 'reorg']
+
+        assert len(data_messages) == 6
+        assert len(reorg_messages) == 1
+
+        assert reorg_messages[0].key == b'reorg:ethereum'
+        assert reorg_messages[0].value['network'] == 'ethereum'
+        assert reorg_messages[0].value['start_block'] == 110
+        assert reorg_messages[0].value['end_block'] == 200
+
+        data_ids = [msg.value['id'] for msg in data_messages]
+        assert data_ids == [1, 2, 3, 4, 5, 6]
+
+    def test_transaction_rollback_on_error(self, kafka_test_config):
+        loader = KafkaLoader(kafka_test_config)
+        topic_name = 'test_rollback_topic'
+
+        batch = pa.RecordBatch.from_pydict(
+            {
+                'id': [1, 2, 3, 4, 5],
+                'name': ['alice', 'bob', 'charlie', 'dave', 'eve'],
+                'value': [100, 200, 300, 400, 500],
+            }
+        )
+
+        with loader:
+            call_count = [0]
+
+            original_send = loader._producer.send
+
+            def failing_send(*args, **kwargs):
+                call_count[0] += 1
+                if call_count[0] == 3:
+                    raise KafkaError('Simulated Kafka send failure')
+                return original_send(*args, **kwargs)
+
+            with patch.object(loader._producer, 'send', side_effect=failing_send):
+                with pytest.raises(RuntimeError, match='FATAL: Permanent error loading batch'):
+                    loader.load_batch(batch, topic_name)
+
+        consumer = KafkaConsumer(
+            topic_name,
+            bootstrap_servers=kafka_test_config['bootstrap_servers'],
+            auto_offset_reset='earliest',
+            consumer_timeout_ms=5000,
+            value_deserializer=lambda m: json.loads(m.decode('utf-8')),
+        )
+
+        messages = list(consumer)
+        consumer.close()
+
+        assert len(messages) == 0