Merge pull request #221 from dpkp/minor_cleanups

Minor cleanups

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),
         )