#!/usr/bin/env python3 """ Unit tests for the BatchProcessor class (_batch_processor.py) This module tests the BatchProcessor class to ensure proper message batching, topic grouping, and future management. """ from confluent_kafka.experimental.aio.producer._kafka_batch_executor import ProducerBatchExecutor as KafkaBatchExecutor from confluent_kafka.experimental.aio.producer._AIOProducer import AIOProducer from confluent_kafka.experimental.aio.producer._producer_batch_processor import ( ProducerBatchManager as ProducerBatchProcessor ) import asyncio import unittest from unittest.mock import Mock, patch import sys import os import concurrent.futures import confluent_kafka # Add src to path for imports sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) class TestProducerBatchProcessor(unittest.TestCase): """Test cases for ProducerBatchProcessor class""" def setUp(self): self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) self.producer_config = { 'bootstrap.servers': 'localhost:9092', 'client.id': 'test-producer', 'message.timeout.ms': 100, 'queue.buffering.max.messages': 1000, 'api.version.request': False, } self.confluent_kafka_producer = confluent_kafka.Producer(self.producer_config) self.kafka_executor = KafkaBatchExecutor(self.confluent_kafka_producer, self.executor) self.batch_processor = ProducerBatchProcessor(self.kafka_executor) async def create_aio_producer(): return AIOProducer(self.producer_config, executor=self.executor) self.aio_producer = self.loop.run_until_complete(create_aio_producer()) def tearDown(self): try: self.loop.run_until_complete(self.aio_producer.close()) except Exception: pass try: self.confluent_kafka_producer.flush(timeout=1) except Exception: pass try: self.executor.shutdown(wait=True, timeout=1) except Exception: pass self.loop.close() def test_basic_functionality(self): self.assertEqual(self.batch_processor.get_buffer_size(), 0) self.assertTrue(self.batch_processor.is_buffer_empty()) future1 = Mock() future2 = Mock() msg_data1 = {'topic': 'topic1', 'value': 'test1', 'key': 'key1'} msg_data2 = {'topic': 'topic2', 'value': 'test2', 'key': 'key2'} self.batch_processor.add_message(msg_data1, future1) self.batch_processor.add_message(msg_data2, future2) self.assertEqual(self.batch_processor.get_buffer_size(), 2) self.assertFalse(self.batch_processor.is_buffer_empty()) self.batch_processor.clear_buffer() self.assertEqual(self.batch_processor.get_buffer_size(), 0) self.assertTrue(self.batch_processor.is_buffer_empty()) def test_group_messages_by_topic(self): future1 = Mock() future2 = Mock() future3 = Mock() msg1 = {'topic': 'topic1', 'value': 'test1', 'user_callback': Mock()} msg2 = {'topic': 'topic2', 'value': 'test2'} msg3 = {'topic': 'topic1', 'value': 'test3', 'user_callback': Mock()} self.batch_processor.add_message(msg1, future1) self.batch_processor.add_message(msg2, future2) self.batch_processor.add_message(msg3, future3) topic_groups = self.batch_processor._group_messages_by_topic_and_partition() self.assertEqual(len(topic_groups), 2) self.assertIn(('topic1', -1), topic_groups) topic1_group = topic_groups[('topic1', -1)] self.assertEqual(len(topic1_group['messages']), 2) self.assertEqual(len(topic1_group['futures']), 2) self.assertEqual(topic1_group['futures'][0], future1) self.assertEqual(topic1_group['futures'][1], future3) self.assertIn(('topic2', -1), topic_groups) topic2_group = topic_groups[('topic2', -1)] self.assertEqual(len(topic2_group['messages']), 1) self.assertEqual(len(topic2_group['futures']), 1) self.assertEqual(topic2_group['futures'][0], future2) def test_prepare_batch_messages(self): messages = [ {'topic': 'test', 'value': 'test1', 'user_callback': Mock(), 'key': 'key1'}, {'topic': 'test', 'value': 'test2', 'partition': 1}, ] batch_messages = self.batch_processor._prepare_batch_messages(messages) self.assertEqual(len(batch_messages), 2) self.assertNotIn('topic', batch_messages[0]) self.assertIn('value', batch_messages[0]) self.assertIn('key', batch_messages[0]) self.assertNotIn('topic', batch_messages[1]) self.assertIn('value', batch_messages[1]) self.assertIn('partition', batch_messages[1]) def test_assign_future_callbacks(self): batch_messages = [ {'value': 'test1'}, {'value': 'test2'}, ] futures = [Mock(), Mock()] self.batch_processor._assign_future_callbacks(batch_messages, futures) self.assertIn('callback', batch_messages[0]) self.assertIn('callback', batch_messages[1]) def test_handle_batch_failure(self): """Test handling batch failures""" futures = [Mock(), Mock()] futures[0].done.return_value = False futures[1].done.return_value = True # Already done exception = RuntimeError("Batch failed") # Handle batch failure self.batch_processor._handle_batch_failure( exception, futures ) # Verify first future got exception (not already done) futures[0].set_exception.assert_called_once_with(exception) # Verify second future was not modified (already done) futures[1].set_exception.assert_not_called() # Note: For real AIOProducer, the user callback is invoked directly by _handle_user_callback def test_flush_empty_buffer(self): """Test flushing empty buffer is no-op""" async def async_test(): await self.batch_processor.flush_buffer() self.assertTrue(self.batch_processor.is_buffer_empty()) self.loop.run_until_complete(async_test()) def test_flush_buffer_with_messages(self): """Test successful buffer flush with messages""" async def async_test(): future1 = self.loop.create_future() future2 = self.loop.create_future() msg1 = {'topic': 'topic1', 'value': 'test1'} msg2 = {'topic': 'topic1', 'value': 'test2'} self.batch_processor.add_message(msg1, future1) self.batch_processor.add_message(msg2, future2) success_future = self.loop.create_future() success_future.set_result("success") with patch.object(self.loop, 'run_in_executor', return_value=success_future): await self.batch_processor.flush_buffer() self.assertTrue(self.batch_processor.is_buffer_empty()) self.loop.run_until_complete(async_test()) def test_flush_buffer_selective_topic(self): """Test selective topic flushing""" async def async_test(): future3 = self.loop.create_future() future4 = self.loop.create_future() msg3 = {'topic': 'topic1', 'value': 'test3'} msg4 = {'topic': 'topic2', 'value': 'test4'} self.batch_processor.add_message(msg3, future3) self.batch_processor.add_message(msg4, future4) success_future = self.loop.create_future() success_future.set_result("success") with patch.object(self.loop, 'run_in_executor', return_value=success_future): await self.batch_processor.flush_buffer(target_topic='topic1') self.assertEqual(self.batch_processor.get_buffer_size(), 1) self.loop.run_until_complete(async_test()) def test_flush_buffer_exception_handling(self): """Test exception handling during buffer flush""" async def async_test(): future = self.loop.create_future() msg = {'topic': 'topic1', 'value': 'test'} self.batch_processor.add_message(msg, future) exception = RuntimeError("Batch execution failed") with patch.object(self.loop, 'run_in_executor', side_effect=exception): with self.assertRaises(RuntimeError): await self.batch_processor.flush_buffer() # Buffer should NOT be empty on exception - we want to retry self.assertFalse(self.batch_processor.is_buffer_empty()) self.assertEqual(self.batch_processor.get_buffer_size(), 1) self.loop.run_until_complete(async_test()) def test_kafka_executor_integration(self): """Test executing a batch operation via KafkaBatchExecutor""" async def async_test(): batch_messages = [ {'value': 'test1', 'callback': Mock()}, {'value': 'test2', 'callback': Mock()}, ] # Mock the executor to return a completed future future_result = self.loop.create_future() future_result.set_result("poll_result") with patch.object(self.loop, 'run_in_executor', return_value=future_result) as mock_run_in_executor: result = await self.kafka_executor.execute_batch('test-topic', batch_messages) # Verify run_in_executor was called mock_run_in_executor.assert_called_once() self.assertEqual(result, "poll_result") self.loop.run_until_complete(async_test()) def _create_mixed_topic_messages(self): """Helper to create messages across multiple topics""" messages_data = [] futures = [] user_callbacks = [] for i in range(4): future = Mock() future.done.return_value = False user_callback = Mock() msg_data = { 'topic': f'topic{i % 2}', 'value': f'unique_value_{i}', 'key': f'unique_key_{i}', 'user_callback': user_callback } self.batch_processor.add_message(msg_data, future) messages_data.append(msg_data) futures.append(future) user_callbacks.append(user_callback) return messages_data, futures, user_callbacks def _add_alternating_topic_messages(self): """Helper to add messages alternating between two topics""" futures = [] for i in range(5): future = self.loop.create_future() msg_data = { 'topic': f'topic{i % 2}', 'value': f'test{i}', 'key': f'key{i}' } self.batch_processor.add_message(msg_data, future) futures.append(future) return futures def test_batch_cycle_buffer_state(self): """Test buffer state during batch cycle""" self._add_alternating_topic_messages() self.assertEqual(self.batch_processor.get_buffer_size(), 5) self.assertFalse(self.batch_processor.is_buffer_empty()) def test_batch_cycle_topic_grouping(self): """Test topic grouping in batch cycle""" self._add_alternating_topic_messages() topic_groups = self.batch_processor._group_messages_by_topic_and_partition() self.assertEqual(len(topic_groups), 2) self.assertIn(('topic0', -1), topic_groups) self.assertIn(('topic1', -1), topic_groups) self.assertEqual(len(topic_groups[('topic0', -1)]['messages']), 3) self.assertEqual(len(topic_groups[('topic1', -1)]['messages']), 2) def test_batch_cycle_message_preparation(self): """Test message preparation in batch cycle""" self._add_alternating_topic_messages() topic_groups = self.batch_processor._group_messages_by_topic_and_partition() batch_messages = self.batch_processor._prepare_batch_messages( topic_groups[('topic0', -1)]['messages'] ) self.assertEqual(len(batch_messages), 3) for batch_msg in batch_messages: self.assertNotIn('topic', batch_msg) self.assertIn('value', batch_msg) self.assertIn('key', batch_msg) def test_batch_message_preparation_with_mixed_sizes(self): """Test batch message preparation with mixed message sizes""" # Create test messages with different sizes messages = [ {'topic': 'test-topic', 'value': 'small message'}, {'topic': 'test-topic', 'value': 'x' * (5 * 1024 * 1024)}, # Large message {'topic': 'test-topic', 'value': 'another small'}, ] futures = [asyncio.Future(), asyncio.Future(), asyncio.Future()] for msg, future in zip(messages, futures): self.batch_processor.add_message(msg, future) topic_groups = self.batch_processor._group_messages_by_topic_and_partition() topic_data = topic_groups[('test-topic', -1)] batch_messages = self.batch_processor._prepare_batch_messages(topic_data['messages']) self.assertEqual(len(batch_messages), 3) large_msg = next((msg for msg in batch_messages if len(str(msg.get('value', ''))) > 1000), None) self.assertIsNotNone(large_msg) def test_future_based_usage_pattern(self): """Test the recommended Future-based usage pattern instead of callbacks.""" # Create test messages without user callbacks messages = [ {'topic': 'test-topic', 'value': 'test1', 'key': 'key1'}, {'topic': 'test-topic', 'value': 'test2', 'key': 'key2'}, ] futures = [asyncio.Future(), asyncio.Future()] # Add messages to batch processor for msg, future in zip(messages, futures): self.batch_processor.add_message(msg, future) # Verify messages are in buffer self.assertEqual(self.batch_processor.get_buffer_size(), 2) # Simulate successful delivery by resolving futures mock_msg1 = Mock() mock_msg1.topic.return_value = 'test-topic' mock_msg1.value.return_value = b'test1' mock_msg2 = Mock() mock_msg2.topic.return_value = 'test-topic' mock_msg2.value.return_value = b'test2' # Applications should await these futures to get delivery results futures[0].set_result(mock_msg1) futures[1].set_result(mock_msg2) # Verify futures are resolved self.assertTrue(futures[0].done()) self.assertTrue(futures[1].done()) self.assertEqual(futures[0].result(), mock_msg1) self.assertEqual(futures[1].result(), mock_msg2) def test_future_based_error_handling(self): """Test Future-based error handling pattern.""" # Create test message message = {'topic': 'test-topic', 'value': 'test', 'key': 'key'} future = asyncio.Future() # Add message to batch processor self.batch_processor.add_message(message, future) # Simulate delivery error by setting exception on future mock_error = RuntimeError("Delivery failed") future.set_exception(mock_error) # Verify future is resolved with exception self.assertTrue(future.done()) with self.assertRaises(RuntimeError): future.result() def test_add_batches_back_to_buffer_basic(self): """Test adding batches back to buffer with basic message data""" from confluent_kafka.experimental.aio.producer._message_batch import create_message_batch # Create test futures future1 = asyncio.Future() future2 = asyncio.Future() # Create test batch with basic message data batch = create_message_batch( topic='test-topic', messages=[ {'value': 'test1', 'key': 'key1'}, {'value': 'test2', 'key': 'key2'} ], futures=[future1, future2], partition=0 ) # Ensure buffer is initially empty self.assertTrue(self.batch_processor.is_buffer_empty()) # Add batch back to buffer self.batch_processor._add_batches_back_to_buffer([batch]) # Verify buffer state self.assertEqual(self.batch_processor.get_buffer_size(), 2) self.assertFalse(self.batch_processor.is_buffer_empty()) # Verify message data was reconstructed correctly self.assertEqual(self.batch_processor._message_buffer[0]['topic'], 'test-topic') self.assertEqual(self.batch_processor._message_buffer[0]['value'], 'test1') self.assertEqual(self.batch_processor._message_buffer[0]['key'], 'key1') self.assertEqual(self.batch_processor._message_buffer[1]['topic'], 'test-topic') self.assertEqual(self.batch_processor._message_buffer[1]['value'], 'test2') self.assertEqual(self.batch_processor._message_buffer[1]['key'], 'key2') # Verify futures are preserved self.assertEqual(self.batch_processor._buffer_futures[0], future1) self.assertEqual(self.batch_processor._buffer_futures[1], future2) def test_add_batches_back_to_buffer_empty_batch(self): """Test adding empty batch back to buffer""" from confluent_kafka.experimental.aio.producer._message_batch import create_message_batch # Create empty batch batch = create_message_batch( topic='test-topic', messages=[], futures=[], partition=0 ) initial_size = self.batch_processor.get_buffer_size() # Add empty batch back self.batch_processor._add_batches_back_to_buffer([batch]) # Buffer size should remain unchanged self.assertEqual(self.batch_processor.get_buffer_size(), initial_size) if __name__ == '__main__': # Run all tests unittest.main(verbosity=2)