def test_send_handler_sets_heal_flag(self):
self.set_responses(
broker_id=1, api="fetch",
responses=[Mock(api="fetch")],
)
self.set_responses(
broker_id=8, api="offset",
responses=[Mock(api="offset")],
)
c = client.Client(["kafka01", "kafka02"])
c.handle_fetch_response = Mock()
c.handle_offset_response = Mock()
def handle_response(response):
c.heal_cluster = True
return "%s handled!" % response.api
c.handle_fetch_response.side_effect = handle_response
c.handle_offset_response.side_effect = handle_response
results = yield c.send({1: Mock(api="fetch"), 8: Mock(api="offset")})
self.assertEqual(
results,
{
1: "fetch handled!",
8: "offset handled!",
}
)
c.cluster.heal.assert_called_once_with()
def test_send_error_on_one_broker(self):
metadata_response = Mock(api="group_coordinator")
self.set_responses(
broker_id=1, api="metadata",
responses=[Exception()],
)
self.set_responses(
broker_id=8, api="group_coordinator",
responses=[metadata_response]
)
c = client.Client(["kafka01", "kafka02"])
c.handle_group_coordinator_response = Mock()
request1 = Mock(api="metadata")
request2 = Mock(api="group_coordinator")
results = yield c.send({1: request1, 8: request2})
self.assertEqual(
results,
{8: c.handle_group_coordinator_response.return_value}
)
self.assert_sent(1, request1)
self.assert_sent(8, request2)
c.cluster.heal.assert_called_once_with()
def test_send_with_async_handlers(self):
self.set_responses(
broker_id=1, api="group_coordinator",
responses=[Mock(api="group_coordinator")],
)
self.set_responses(
broker_id=8, api="fetch",
responses=[Mock(api="fetch")],
)
c = client.Client(["kafka01", "kafka02"])
c.handle_group_coordinator_response = Mock()
c.handle_group_coordinator_response.return_value = self.future_value(
"metadata handled!"
)
c.handle_fetch_response = Mock()
c.handle_fetch_response.return_value = self.future_value(
"fetch handled!"
)
results = yield c.send(
{1: Mock(api="group_coordinator"), 8: Mock(api="fetch")}
)
self.assertEqual(
results,
{
1: "metadata handled!",
8: "fetch handled!",
}
)
self.assertFalse(c.cluster.heal.called)
def test_close(self):
c = client.Client([])
c.wind_down = Mock()
c.wind_down.return_value = self.future_value(None)
yield c.close()
self.assertEqual(c.closing, True)
c.wind_down.assert_called_once_with()
c.cluster.stop.assert_called_once_with()
def test_wind_down_must_be_implemented(self):
c = client.Client([])
error = None
try:
yield c.wind_down()
except Exception as e:
error = e
self.assertIsInstance(error, NotImplementedError)
def test_send_no_handler(self):
self.set_responses(
broker_id=8, api="produce",
responses=[Mock(api="produce")],
)
c = client.Client(["kafka01", "kafka02"])
error = None
try:
yield c.send({8: Mock(api="produce")})
except Exception as e:
error = e
self.assertIsInstance(error, exc.UnhandledResponseError)
self.assertEqual(error.api, "produce")
def test_send_stream_closed(self):
self.set_responses(
broker_id=1, api="metadata",
responses=[
iostream.StreamClosedError(),
]
)
c = client.Client(["kafka01", "kafka02"])
request = Mock(api="metadata")
results = yield c.send({1: request})
self.assert_sent(1, request)
self.assertEqual(c.cluster.heal.called, False)
self.assertEqual(results, {})
def test_send_connection_error(self):
self.set_responses(
broker_id=1, api="metadata",
responses=[
exc.BrokerConnectionError("kafka01", 1234)
]
)
c = client.Client(["kafka01", "kafka02"])
request = Mock(api="metadata")
results = yield c.send({1: request})
self.assert_sent(1, request)
c.cluster.heal.assert_called_once_with()
self.assertEqual(results, {})
def test_send_with_handlers(self):
metadata_response = Mock(api="group_coordinator")
fetch_response = Mock(api="fetch")
self.set_responses(
broker_id=1, api="fetch",
responses=[fetch_response],
)
self.set_responses(
broker_id=8, api="group_coordinator",
responses=[metadata_response]
)
c = client.Client(["kafka01", "kafka02"])
c.handle_group_coordinator_response = Mock()
c.handle_fetch_response = Mock()
request1 = Mock(api="fetch")
request2 = Mock(api="group_coordinator")
results = yield c.send({1: request1, 8: request2})
self.assertEqual(
results,
{
1: c.handle_fetch_response.return_value,
8: c.handle_group_coordinator_response.return_value
}
)
self.assert_sent(1, request1)
self.assert_sent(8, request2)
c.handle_group_coordinator_response.assert_called_once_with(
metadata_response
)
c.handle_fetch_response.assert_called_once_with(fetch_response)
self.assertFalse(c.cluster.heal.called)
def test_connect_starts_cluster(self):
c = client.Client([])
yield c.connect()
c.cluster.start.assert_called_once_with()
def test_instantiation(self):
c = client.Client(["kafka01", "kafka02"])
self.assertEqual(c.closing, False)
self.assertEqual(c.heal_cluster, False)