diff options
| -rw-r--r-- | kafka/protocol/legacy.py | 2 | ||||
| -rw-r--r-- | test/test_protocol.py | 6 |
2 files changed, 5 insertions, 3 deletions
diff --git a/kafka/protocol/legacy.py b/kafka/protocol/legacy.py index 412a957..d7ac50a 100644 --- a/kafka/protocol/legacy.py +++ b/kafka/protocol/legacy.py @@ -1,3 +1,5 @@ +from __future__ import absolute_import + import logging import struct diff --git a/test/test_protocol.py b/test/test_protocol.py index ac7bea6..368c2d0 100644 --- a/test/test_protocol.py +++ b/test/test_protocol.py @@ -780,11 +780,11 @@ class TestProtocol(unittest.TestCase): @contextmanager def mock_create_message_fns(self): import kafka.protocol - with patch.object(kafka.protocol, "create_message", + with patch.object(kafka.protocol.legacy, "create_message", return_value=sentinel.message): - with patch.object(kafka.protocol, "create_gzip_message", + with patch.object(kafka.protocol.legacy, "create_gzip_message", return_value=sentinel.gzip_message): - with patch.object(kafka.protocol, "create_snappy_message", + with patch.object(kafka.protocol.legacy, "create_snappy_message", return_value=sentinel.snappy_message): yield |