1 files changed, 10 insertions, 6 deletions
diff --git a/test/test_producer.py b/test/test_producer.py
index 119930f..f6b3d6a 100644
--- a/test/test_producer.py
+++ b/test/test_producer.py
@@ -7,6 +7,7 @@ from . import unittest
 
 from kafka.producer.base import Producer
 
+
 class TestKafkaProducer(unittest.TestCase):
     def test_producer_message_types(self):
 
@@ -28,11 +29,14 @@ class TestKafkaProducer(unittest.TestCase):
     def test_topic_message_types(self):
         from kafka.producer.simple import SimpleProducer
 
-        producer = SimpleProducer(MagicMock())
-        topic = "test-topic"
-        partition = 0
+        client = MagicMock()
 
-        def send_message():
-            producer.send_messages(topic, partition, b'hi')
+        def partitions(topic):
+            return [0, 1]
 
-        self.assertRaises(TypeError, send_message)
+        client.get_partition_ids_for_topic = partitions
+
+        producer = SimpleProducer(client, random_start=False)
+        topic = b"test-topic"
+        producer.send_messages(topic, b'hi')
+        assert client.send_produce_request.called