From b99a9eee64aa9dec34f43085a551d076513b9f51 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Sun, 31 Aug 2014 11:13:02 -0700 Subject: Import cleanups in tests --- test/test_protocol.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) (limited to 'test/test_protocol.py') 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), ) -- cgit v1.2.1