def test__send_sasl_token(self):
# Before Kafka 1.0.0 SASL was performed on the wire without
# KAFKA_HEADER in the protocol. So we needed another private
# function to send `raw` data with only length prefixed
# setup connection with mocked transport and protocol
conn = AIOKafkaConnection(
host="", port=9999, loop=self.loop
)
conn.close = mock.MagicMock()
conn._writer = mock.MagicMock()
out_buffer = []
conn._writer.write = mock.Mock(side_effect=out_buffer.append)
conn._reader = mock.MagicMock()
self.assertEqual(len(conn._requests), 0)
# Successful send
fut = conn._send_sasl_token(b"Super data")
self.assertEqual(b''.join(out_buffer), b"\x00\x00\x00\nSuper data")
self.assertEqual(len(conn._requests), 1)
out_buffer.clear()
# Resolve the request
conn._requests[0][2].set_result(None)
conn._requests.clear()
yield from fut
# Broken pipe error
conn._writer.write.side_effect = OSError
with self.assertRaises(ConnectionError):
conn._send_sasl_token(b"Super data")
self.assertEqual(out_buffer, [])
self.assertEqual(len(conn._requests), 0)
self.assertEqual(conn.close.call_count, 1)
conn._writer = None
with self.assertRaises(ConnectionError):
conn._send_sasl_token(b"Super data")
# We don't need to close 2ce
self.assertEqual(conn.close.call_count, 1)
async def test__send_sasl_token(self):
# Before Kafka 1.0.0 SASL was performed on the wire without
# KAFKA_HEADER in the protocol. So we needed another private
# function to send `raw` data with only length prefixed
# setup connection with mocked transport and protocol
conn = AIOKafkaConnection(
host="", port=9999, loop=self.loop
)
conn.close = mock.MagicMock()
conn._writer = mock.MagicMock()
out_buffer = []
conn._writer.write = mock.Mock(side_effect=out_buffer.append)
conn._reader = mock.MagicMock()
self.assertEqual(len(conn._requests), 0)
# Successful send
fut = conn._send_sasl_token(b"Super data")
self.assertEqual(b''.join(out_buffer), b"\x00\x00\x00\nSuper data")
self.assertEqual(len(conn._requests), 1)
out_buffer.clear()
# Resolve the request
conn._requests[0][2].set_result(None)
conn._requests.clear()
await fut
# Broken pipe error
conn._writer.write.side_effect = OSError
with self.assertRaises(ConnectionError):
conn._send_sasl_token(b"Super data")
self.assertEqual(out_buffer, [])
self.assertEqual(len(conn._requests), 0)
self.assertEqual(conn.close.call_count, 1)
conn._writer = None
with self.assertRaises(ConnectionError):
conn._send_sasl_token(b"Super data")
# We don't need to close 2ce
self.assertEqual(conn.close.call_count, 1)
def test__do_sasl_handshake_v0(self):
host, port = self.kafka_host, self.kafka_port
# setup connection with mocked send and send_bytes
conn = AIOKafkaConnection(
host=host, port=port, loop=self.loop,
sasl_mechanism="PLAIN",
sasl_plain_username="******",
sasl_plain_password="******"
)
conn.close = close_mock = mock.MagicMock()
supported_mechanisms = ["PLAIN"]
error_class = NoError
@asyncio.coroutine
def mock_send(request, expect_response=True):
return SaslHandShakeResponse[0](
error_code=error_class.errno,
enabled_mechanisms=supported_mechanisms
)
@asyncio.coroutine
def mock_sasl_send(payload, expect_response):
return b""
conn.send = mock.Mock(side_effect=mock_send)
conn._send_sasl_token = mock.Mock(side_effect=mock_sasl_send)
conn._version_info = VersionInfo({
SaslHandShakeRequest[0].API_KEY: [0, 0]
})
yield from conn._do_sasl_handshake()
supported_mechanisms = ["GSSAPI"]
with self.assertRaises(UnsupportedSaslMechanismError):
yield from conn._do_sasl_handshake()
self.assertTrue(close_mock.call_count)
error_class = UnknownError
close_mock.reset()
with self.assertRaises(UnknownError):
yield from conn._do_sasl_handshake()
self.assertTrue(close_mock.call_count)
def test__do_sasl_handshake_v0(self):
host, port = self.kafka_host, self.kafka_port
# setup connection with mocked send and send_bytes
conn = AIOKafkaConnection(host=host,
port=port,
loop=self.loop,
sasl_mechanism="PLAIN",
sasl_plain_username="******",
sasl_plain_password="******")
conn.close = close_mock = mock.MagicMock()
supported_mechanisms = ["PLAIN"]
error_class = NoError
@asyncio.coroutine
def mock_send(request, expect_response=True):
return SaslHandShakeResponse[0](
error_code=error_class.errno,
enabled_mechanisms=supported_mechanisms)
@asyncio.coroutine
def mock_sasl_send(payload, expect_response):
return b""
conn.send = mock.Mock(side_effect=mock_send)
conn._send_sasl_token = mock.Mock(side_effect=mock_sasl_send)
conn._version_info = VersionInfo(
{SaslHandShakeRequest[0].API_KEY: [0, 0]})
yield from conn._do_sasl_handshake()
supported_mechanisms = ["GSSAPI"]
with self.assertRaises(UnsupportedSaslMechanismError):
yield from conn._do_sasl_handshake()
self.assertTrue(close_mock.call_count)
error_class = UnknownError
close_mock.reset()
with self.assertRaises(UnknownError):
yield from conn._do_sasl_handshake()
self.assertTrue(close_mock.call_count)
