diff options
| author | Dana Powers <dana.powers@rd.io> | 2014-08-31 11:13:02 -0700 |
|---|---|---|
| committer | Dana Powers <dana.powers@rd.io> | 2014-09-04 18:05:05 -0700 |
| commit | b99a9eee64aa9dec34f43085a551d076513b9f51 (patch) | |
| tree | 91dbc62e3863897213af2fbd0070b92bff720e28 /test/test_protocol.py | |
| parent | 404c56a55d95783a027313ae69d43097ceb18ab6 (diff) | |
| download | kafka-python-b99a9eee64aa9dec34f43085a551d076513b9f51.tar.gz | |
Import cleanups in tests
Diffstat (limited to 'test/test_protocol.py')
| -rw-r--r-- | test/test_protocol.py | 16 |
1 files changed, 7 insertions, 9 deletions
diff --git a/test/test_protocol.py b/test/test_protocol.py index 2089f48..b1c5201 100644 --- a/test/test_protocol.py +++ b/test/test_protocol.py @@ -1,10 +1,8 @@ -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 ( @@ -701,12 +699,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), ) |