confluent-kafka-python/tests/test_kafka_batch_executor.py at master · confluentinc/confluent-kafka-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
#!/usr/bin/env python3
"""
Unit tests for the KafkaBatchExecutor class (_kafka_batch_executor.py)

This module tests the KafkaBatchExecutor class to ensure proper
Kafka batch execution and partial failure handling.
"""

import asyncio
import concurrent.futures
import os
import sys
import unittest
from unittest.mock import Mock, patch

import confluent_kafka
from confluent_kafka.aio.producer._kafka_batch_executor import ProducerBatchExecutor as KafkaBatchExecutor

# Add src to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))


class TestKafkaBatchExecutor(unittest.TestCase):

    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
        self.mock_producer = Mock(spec=confluent_kafka.Producer)
        self.kafka_executor = KafkaBatchExecutor(self.mock_producer, self.executor)

    def tearDown(self):
        self.executor.shutdown(wait=True)
        self.loop.close()

    def test_initialization(self):
        self.assertEqual(self.kafka_executor._producer, self.mock_producer)
        self.assertEqual(self.kafka_executor._executor, self.executor)

    def test_execute_batch_success(self):
        async def async_test():
            batch_messages = [
                {'value': 'test1', 'callback': Mock()},
                {'value': 'test2', 'callback': Mock()},
            ]

            self.mock_producer.produce_batch.return_value = None
            self.mock_producer.poll.return_value = 2

            result = await self.kafka_executor.execute_batch('test-topic', batch_messages)

            self.mock_producer.produce_batch.assert_called_once_with('test-topic', batch_messages, partition=-1)
            self.mock_producer.poll.assert_called_once_with(0)
            self.assertEqual(result, 2)

        self.loop.run_until_complete(async_test())

    def test_partial_failure_handling(self):
        async def async_test():
            callback1 = Mock()
            callback2 = Mock()
            batch_messages = [
                {'value': 'test1', 'callback': callback1},
                {'value': 'test2', 'callback': callback2, '_error': 'MSG_SIZE_TOO_LARGE'},
            ]

            self.mock_producer.produce_batch.return_value = None
            self.mock_producer.poll.return_value = 1

            result = await self.kafka_executor.execute_batch('test-topic', batch_messages)

            self.mock_producer.produce_batch.assert_called_once_with('test-topic', batch_messages, partition=-1)
            self.mock_producer.poll.assert_called_once_with(0)

            callback1.assert_not_called()
            callback2.assert_called_once_with('MSG_SIZE_TOO_LARGE', None)

            self.assertEqual(result, 1)

        self.loop.run_until_complete(async_test())

    def test_batch_execution_exception(self):
        async def async_test():
            batch_messages = [{'value': 'test1', 'callback': Mock()}]

            self.mock_producer.produce_batch.side_effect = Exception("Kafka error")

            with self.assertRaises(Exception) as context:
                await self.kafka_executor.execute_batch('test-topic', batch_messages)

            self.assertEqual(str(context.exception), "Kafka error")
            self.mock_producer.produce_batch.assert_called_once_with('test-topic', batch_messages, partition=-1)

        self.loop.run_until_complete(async_test())

    def test_callback_exception_handling(self):
        async def async_test():
            failing_callback = Mock(side_effect=Exception("Callback error"))

            batch_messages = [
                {'value': 'test1', 'callback': failing_callback, '_error': 'TEST_ERROR'},
            ]

            self.mock_producer.produce_batch.return_value = None
            self.mock_producer.poll.return_value = 0

            # Expect the callback exception to be raised
            with self.assertRaises(Exception) as context:
                await self.kafka_executor.execute_batch('test-topic', batch_messages)

            # Verify the callback was called before the exception
            failing_callback.assert_called_once_with('TEST_ERROR', None)
            self.assertEqual(str(context.exception), "Callback error")

        self.loop.run_until_complete(async_test())

    def test_thread_pool_execution(self):
        async def async_test():
            batch_messages = [{'value': 'test1', 'callback': Mock()}]

            with patch.object(self.loop, 'run_in_executor') as mock_run_in_executor:
                future_result = self.loop.create_future()
                future_result.set_result(1)
                mock_run_in_executor.return_value = future_result

                result = await self.kafka_executor.execute_batch('test-topic', batch_messages)

                mock_run_in_executor.assert_called_once()
                args = mock_run_in_executor.call_args
                self.assertEqual(args[0][0], self.executor)
                self.assertTrue(callable(args[0][1]))

                self.assertEqual(result, 1)

        self.loop.run_until_complete(async_test())


if __name__ == '__main__':
    unittest.main()