Skip to content

Commit 220c09f

Browse files
mattoberledpkp
authored andcommitted
Tests AWS_MSK_IAM signature and payload generation
The two tests in `test/test_msk.py` should ensure that the changes to `kafka/msk.py` do not break the authentication payload. The authentication payload was validated using a real AWS Kafka cluster before adding tests with the hard-coded signatures.
1 parent 8e47dc0 commit 220c09f

File tree

1 file changed

+67
-0
lines changed

1 file changed

+67
-0
lines changed

test/test_msk.py

+67
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import datetime
2+
import json
3+
4+
from kafka.msk import AwsMskIamClient
5+
6+
try:
7+
from unittest import mock
8+
except ImportError:
9+
import mock
10+
11+
12+
def client_factory(token=None):
13+
now = datetime.datetime.utcfromtimestamp(1629321911)
14+
with mock.patch('kafka.msk.datetime') as mock_dt:
15+
mock_dt.datetime.utcnow = mock.Mock(return_value=now)
16+
return AwsMskIamClient(
17+
host='localhost',
18+
access_key='XXXXXXXXXXXXXXXXXXXX',
19+
secret_key='XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX',
20+
region='us-east-1',
21+
token=token,
22+
)
23+
24+
25+
def test_aws_msk_iam_client_permanent_credentials():
26+
client = client_factory(token=None)
27+
msg = client.first_message()
28+
assert msg
29+
assert isinstance(msg, bytes)
30+
actual = json.loads(msg)
31+
32+
expected = {
33+
'version': '2020_10_22',
34+
'host': 'localhost',
35+
'user-agent': 'kafka-python',
36+
'action': 'kafka-cluster:Connect',
37+
'x-amz-algorithm': 'AWS4-HMAC-SHA256',
38+
'x-amz-credential': 'XXXXXXXXXXXXXXXXXXXX/20210818/us-east-1/kafka-cluster/aws4_request',
39+
'x-amz-date': '20210818T212511Z',
40+
'x-amz-signedheaders': 'host',
41+
'x-amz-expires': '900',
42+
'x-amz-signature': '0fa42ae3d5693777942a7a4028b564f0b372bafa2f71c1a19ad60680e6cb994b',
43+
}
44+
assert actual == expected
45+
46+
47+
def test_aws_msk_iam_client_temporary_credentials():
48+
client = client_factory(token='XXXXX')
49+
msg = client.first_message()
50+
assert msg
51+
assert isinstance(msg, bytes)
52+
actual = json.loads(msg)
53+
54+
expected = {
55+
'version': '2020_10_22',
56+
'host': 'localhost',
57+
'user-agent': 'kafka-python',
58+
'action': 'kafka-cluster:Connect',
59+
'x-amz-algorithm': 'AWS4-HMAC-SHA256',
60+
'x-amz-credential': 'XXXXXXXXXXXXXXXXXXXX/20210818/us-east-1/kafka-cluster/aws4_request',
61+
'x-amz-date': '20210818T212511Z',
62+
'x-amz-signedheaders': 'host',
63+
'x-amz-expires': '900',
64+
'x-amz-signature': 'b0619c50b7ecb4a7f6f92bd5f733770df5710e97b25146f97015c0b1db783b05',
65+
'x-amz-security-token': 'XXXXX',
66+
}
67+
assert actual == expected

0 commit comments

Comments
 (0)