-
Notifications
You must be signed in to change notification settings - Fork 932
Add _len_ function to AIOProducer #2140
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
k-raina
wants to merge
2
commits into
master
Choose a base branch
from
add-len-to-AIOProducer
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+118
−0
Open
Changes from all commits
Commits
Show all changes
2 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -534,3 +534,100 @@ async def test_edge_cases_batching(self, mock_producer, mock_common, basic_confi | |
| assert mock_flush.call_count >= 1 # At least one flush | ||
|
|
||
| await producer.close() | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_aio_producer_len_with_buffered_messages(self, mock_producer, mock_common, basic_config): | ||
| """Test that __len__ counts messages in async batch buffer""" | ||
| producer = AIOProducer(basic_config, batch_size=10, buffer_timeout=0) | ||
|
|
||
| # Produce 5 messages (less than batch_size, so they stay in buffer) | ||
| with patch.object(producer, '_flush_buffer') as mock_flush: # Prevent auto-flush | ||
| for i in range(5): | ||
| await producer.produce('test-topic', value=f'msg-{i}'.encode()) | ||
|
|
||
| # Verify flush was not called (messages stayed in buffer) | ||
| mock_flush.assert_not_called() | ||
|
|
||
| # len() should count messages in buffer | ||
| assert len(producer) == 5 | ||
| assert producer._batch_processor.get_buffer_size() == 5 | ||
| assert len(producer._producer) == 0 # Nothing in librdkafka yet | ||
| # Verify len() equals the sum | ||
| assert len(producer) == producer._batch_processor.get_buffer_size() + len(producer._producer) | ||
|
|
||
| await producer.close() | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_aio_producer_len_after_flush(self, mock_producer, mock_common, basic_config): | ||
| """Test that __len__ counts messages after flush to librdkafka""" | ||
| producer = AIOProducer(basic_config, batch_size=10, buffer_timeout=0) | ||
|
|
||
| # Produce and flush | ||
| with patch.object(producer, '_flush_buffer'): # Prevent auto-flush | ||
| for i in range(5): | ||
| await producer.produce('test-topic', value=f'msg-{i}'.encode()) | ||
|
|
||
| # Flush to move messages to librdkafka | ||
| await producer.flush() | ||
|
|
||
| # After flush, messages move to librdkafka queue | ||
| # Verify len() correctly equals the sum of buffer + librdkafka | ||
| buffer_count = producer._batch_processor.get_buffer_size() | ||
| librdkafka_count = len(producer._producer) | ||
| total_len = len(producer) | ||
|
|
||
| # Buffer should be empty after flush | ||
| assert buffer_count == 0 | ||
| # Verify len() equals the sum (this validates the __len__ implementation) | ||
| assert total_len == buffer_count + librdkafka_count | ||
| # Messages should be in librdkafka queue after flush | ||
| assert total_len == librdkafka_count | ||
|
|
||
| await producer.close() | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_aio_producer_len_closed_producer(self, mock_producer, mock_common, basic_config): | ||
| """Test that __len__ returns 0 for closed producer""" | ||
| producer = AIOProducer(basic_config, batch_size=10, buffer_timeout=0) | ||
|
|
||
| # Produce some messages | ||
| with patch.object(producer, '_flush_buffer'): # Prevent auto-flush | ||
| for i in range(3): | ||
| await producer.produce('test-topic', value=f'msg-{i}'.encode()) | ||
|
|
||
| # Verify messages are there | ||
| assert len(producer) == 3 | ||
|
|
||
| # Close producer | ||
| await producer.close() | ||
|
|
||
| # len() should return 0 for closed producer | ||
| assert len(producer) == 0 | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_aio_producer_len_mixed_state(self, mock_producer, mock_common, basic_config): | ||
| """Test __len__ when messages are in both buffer and librdkafka queue""" | ||
| producer = AIOProducer(basic_config, batch_size=5, buffer_timeout=0) | ||
|
|
||
| # Produce 7 messages - first 5 should flush (batch_size=5), last 2 stay in buffer | ||
| with patch.object(producer, '_flush_buffer') as mock_flush: | ||
| for i in range(7): | ||
| await producer.produce('test-topic', value=f'msg-{i}'.encode()) | ||
|
Comment on lines
+613
to
+615
|
||
|
|
||
| # With batch_size=5, flush should be called after 5th message | ||
| # Verify flush was called (at least once when batch_size reached) | ||
| assert mock_flush.call_count >= 1 | ||
|
|
||
| # After batch_size messages, some may have flushed | ||
| # Total should be sum of buffer + librdkafka queue | ||
| buffer_count = producer._batch_processor.get_buffer_size() | ||
| librdkafka_count = len(producer._producer) | ||
| total_count = len(producer) | ||
|
|
||
| # Verify len() correctly equals the sum (this validates the __len__ implementation) | ||
| assert total_count == buffer_count + librdkafka_count | ||
| # At least the messages beyond batch_size should be in buffer | ||
| # (exact count depends on flush behavior) | ||
| assert total_count >= 2 # At least the last 2 should be pending | ||
|
|
||
| await producer.close() | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The test patches
_flush_bufferto prevent auto-flush, but doesn't verify that the patch was actually called or not called. Consider adding an assertion likemock_flush.assert_not_called()after producing messages to ensure the buffer wasn't inadvertently flushed.