diff options
Diffstat (limited to 'test/test_protocol.py')
| -rw-r--r-- | test/test_protocol.py | 25 |
1 files changed, 9 insertions, 16 deletions
diff --git a/test/test_protocol.py b/test/test_protocol.py index 2089f48..f6e3c96 100644 --- a/test/test_protocol.py +++ b/test/test_protocol.py @@ -1,25 +1,18 @@ -import contextlib -from contextlib import contextmanager +from contextlib import contextmanager, nested import struct -import unittest2 -import mock -from mock import sentinel +import unittest2 +from mock import patch, sentinel -from kafka import KafkaClient from kafka.common import ( OffsetRequest, OffsetCommitRequest, OffsetFetchRequest, OffsetResponse, OffsetCommitResponse, OffsetFetchResponse, ProduceRequest, FetchRequest, Message, ChecksumError, ConsumerFetchSizeTooSmall, ProduceResponse, FetchResponse, OffsetAndMessage, - BrokerMetadata, PartitionMetadata, TopicAndPartition, KafkaUnavailableError, - ProtocolError, LeaderUnavailableError, PartitionUnavailableError, + BrokerMetadata, PartitionMetadata, ProtocolError, UnsupportedCodecError ) -from kafka.codec import ( - has_snappy, gzip_encode, gzip_decode, - snappy_encode, snappy_decode -) +from kafka.codec import has_snappy, gzip_decode, snappy_decode import kafka.protocol from kafka.protocol import ( ATTRIBUTE_CODEC_MASK, CODEC_NONE, CODEC_GZIP, CODEC_SNAPPY, KafkaProtocol, @@ -701,12 +694,12 @@ class TestProtocol(unittest2.TestCase): @contextmanager def mock_create_message_fns(self): - patches = contextlib.nested( - mock.patch.object(kafka.protocol, "create_message", + patches = nested( + patch.object(kafka.protocol, "create_message", return_value=sentinel.message), - mock.patch.object(kafka.protocol, "create_gzip_message", + patch.object(kafka.protocol, "create_gzip_message", return_value=sentinel.gzip_message), - mock.patch.object(kafka.protocol, "create_snappy_message", + patch.object(kafka.protocol, "create_snappy_message", return_value=sentinel.snappy_message), ) |