From 7ff323727d99e0c33a68423300e7f88a9cf3f830 Mon Sep 17 00:00:00 2001
From: Matt Oberle <mattoberle@users.noreply.github.com>
Date: Wed, 18 Aug 2021 18:10:35 -0400
Subject: [PATCH] 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.
---
 test/test_msk.py | 67 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 test/test_msk.py

diff --git a/test/test_msk.py b/test/test_msk.py
new file mode 100644
index 000000000..d4620e39e
--- /dev/null
+++ b/test/test_msk.py
@@ -0,0 +1,67 @@
+import datetime
+import json
+
+from kafka.msk import AwsMskIamClient
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+
+def client_factory(token=None):
+    now = datetime.datetime.utcfromtimestamp(1629321911)
+    with mock.patch('kafka.msk.datetime') as mock_dt:
+        mock_dt.datetime.utcnow = mock.Mock(return_value=now)
+        return AwsMskIamClient(
+            host='localhost',
+            access_key='XXXXXXXXXXXXXXXXXXXX',
+            secret_key='XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX',
+            region='us-east-1',
+            token=token,
+        )
+
+
+def test_aws_msk_iam_client_permanent_credentials():
+    client = client_factory(token=None)
+    msg = client.first_message()
+    assert msg
+    assert isinstance(msg, bytes)
+    actual = json.loads(msg)
+
+    expected = {
+        'version': '2020_10_22',
+        'host': 'localhost',
+        'user-agent': 'kafka-python',
+        'action': 'kafka-cluster:Connect',
+        'x-amz-algorithm': 'AWS4-HMAC-SHA256',
+        'x-amz-credential': 'XXXXXXXXXXXXXXXXXXXX/20210818/us-east-1/kafka-cluster/aws4_request',
+        'x-amz-date': '20210818T212511Z',
+        'x-amz-signedheaders': 'host',
+        'x-amz-expires': '900',
+        'x-amz-signature': '0fa42ae3d5693777942a7a4028b564f0b372bafa2f71c1a19ad60680e6cb994b',
+    }
+    assert actual == expected
+
+
+def test_aws_msk_iam_client_temporary_credentials():
+    client = client_factory(token='XXXXX')
+    msg = client.first_message()
+    assert msg
+    assert isinstance(msg, bytes)
+    actual = json.loads(msg)
+
+    expected = {
+        'version': '2020_10_22',
+        'host': 'localhost',
+        'user-agent': 'kafka-python',
+        'action': 'kafka-cluster:Connect',
+        'x-amz-algorithm': 'AWS4-HMAC-SHA256',
+        'x-amz-credential': 'XXXXXXXXXXXXXXXXXXXX/20210818/us-east-1/kafka-cluster/aws4_request',
+        'x-amz-date': '20210818T212511Z',
+        'x-amz-signedheaders': 'host',
+        'x-amz-expires': '900',
+        'x-amz-signature': 'b0619c50b7ecb4a7f6f92bd5f733770df5710e97b25146f97015c0b1db783b05',
+        'x-amz-security-token': 'XXXXX',
+    }
+    assert actual == expected