Fix tests broken by legacy module move

This commit is contained in:
Dana Powers
2015-11-28 19:29:32 +08:00
committed by Zack Dever
parent 3f65ff4ab9
commit dc94b5fe9f
2 changed files with 5 additions and 3 deletions

View File

@@ -1,3 +1,5 @@
from __future__ import absolute_import
import logging
import struct

View File

@@ -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