def test_connection_version_info(self):
# All version supported
version_info = VersionInfo({
SaslHandShakeRequest[0].API_KEY: [0, 1]
})
self.assertEqual(
version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[1])
# Broker only supports the lesser version
version_info = VersionInfo({
SaslHandShakeRequest[0].API_KEY: [0, 0]
})
self.assertEqual(
version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[0])
# We don't support any version compatible with the broker
version_info = VersionInfo({
SaslHandShakeRequest[0].API_KEY: [2, 3]
})
with self.assertRaises(KafkaError):
self.assertEqual(
version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[1])
# No information on the supported versions
version_info = VersionInfo({})
self.assertEqual(
version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[0])
async def test__do_sasl_handshake_v1(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="******",
security_protocol="SASL_PLAINTEXT"
)
conn.close = close_mock = mock.MagicMock()
supported_mechanisms = ["PLAIN"]
error_class = NoError
auth_error_class = NoError
async def mock_send(request, expect_response=True):
if request.API_KEY == SaslHandShakeRequest[0].API_KEY:
assert request.API_VERSION == 1
return SaslHandShakeResponse[1](
error_code=error_class.errno,
enabled_mechanisms=supported_mechanisms
)
else:
assert request.API_KEY == SaslAuthenticateRequest[0].API_KEY
return SaslAuthenticateResponse[0](
error_code=auth_error_class.errno,
error_message="",
sasl_auth_bytes=b""
)
conn.send = mock.Mock(side_effect=mock_send)
conn._version_info = VersionInfo({
SaslHandShakeRequest[0].API_KEY: [0, 1]
})
await conn._do_sasl_handshake()
supported_mechanisms = ["GSSAPI"]
with self.assertRaises(UnsupportedSaslMechanismError):
await conn._do_sasl_handshake()
self.assertTrue(close_mock.call_count)
supported_mechanisms = ["PLAIN"]
auth_error_class = IllegalSaslStateError
close_mock.reset()
with self.assertRaises(IllegalSaslStateError):
await conn._do_sasl_handshake()
self.assertTrue(close_mock.call_count)
auth_error_class = NoError
error_class = UnknownError
close_mock.reset()
with self.assertRaises(UnknownError):
await 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)
def test_connection_version_info(self):
# All version supported
version_info = VersionInfo({SaslHandShakeRequest[0].API_KEY: [0, 1]})
self.assertEqual(version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[1])
# Broker only supports the lesser version
version_info = VersionInfo({SaslHandShakeRequest[0].API_KEY: [0, 0]})
self.assertEqual(version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[0])
# We don't support any version compatible with the broker
version_info = VersionInfo({SaslHandShakeRequest[0].API_KEY: [2, 3]})
with self.assertRaises(KafkaError):
self.assertEqual(version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[1])
# No information on the supported versions
version_info = VersionInfo({})
self.assertEqual(version_info.pick_best(SaslHandShakeRequest),
SaslHandShakeRequest[0])