Repository: dpkp/kafka-python
Pull request #2522 (2025-03-11 23:49:27)

Fix lint issues via ruff check
`ruff check --exclude kafka/vendor/ kafka/ test/`


Patch:
diff --git a/kafka/client_async.py b/kafka/client_async.py
index bd34c3b..b72c05d 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -19,7 +19,7 @@ except ImportError:
 from kafka.vendor import six
 
 from kafka.cluster import ClusterMetadata
-from kafka.conn import BrokerConnection, ConnectionStates, collect_hosts, get_ip_port_afi
+from kafka.conn import BrokerConnection, ConnectionStates, get_ip_port_afi
 from kafka import errors as Errors
 from kafka.future import Future
 from kafka.metrics import AnonMeasurable
diff --git a/kafka/conn.py b/kafka/conn.py
index 857b13a..7287a48 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -14,7 +14,6 @@ except ImportError:
     from kafka.vendor import selectors34 as selectors
 
 import socket
-import struct
 import threading
 import time
 
@@ -23,7 +22,6 @@ from kafka.vendor import six
 import kafka.errors as Errors
 from kafka.future import Future
 from kafka.metrics.stats import Avg, Count, Max, Rate
-from kafka.oauth.abstract import AbstractTokenProvider
 from kafka.protocol.admin import DescribeAclsRequest, DescribeClientQuotasRequest, ListGroupsRequest
 from kafka.protocol.api_versions import ApiVersionsRequest
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
@@ -36,7 +34,7 @@ from kafka.protocol.parser import KafkaProtocol
 from kafka.protocol.produce import ProduceRequest
 from kafka.protocol.sasl_authenticate import SaslAuthenticateRequest
 from kafka.protocol.sasl_handshake import SaslHandshakeRequest
-from kafka.protocol.types import Int32, Int8
+from kafka.protocol.types import Int32
 from kafka.sasl import get_sasl_mechanism
 from kafka.version import __version__
 
@@ -1151,7 +1149,8 @@ class BrokerConnection(object):
     def next_ifr_request_timeout_ms(self):
         with self._lock:
             if self.in_flight_requests:
-                get_timeout = lambda v: v[2]
+                def get_timeout(v):
+                    return v[2]
                 next_timeout = min(map(get_timeout,
                                    self.in_flight_requests.values()))
                 return max(0, (next_timeout - time.time()) * 1000)
@@ -1159,11 +1158,11 @@ class BrokerConnection(object):
                 return float('inf')
 
     def get_api_versions(self):
-        if self._api_versions is not None:
-            return self._api_versions
-
-        version = self.check_version()
-        # _api_versions is set as a side effect of check_versions()
+        # _api_versions is set as a side effect of first connection
+        # which should typically be bootstrap, but call check_version
+        # if that hasn't happened yet
+        if self._api_versions is None:
+            self.check_version()
         return self._api_versions
 
     def _infer_broker_version_from_api_versions(self, api_versions):
@@ -1201,11 +1200,11 @@ class BrokerConnection(object):
         ]
 
         # Get the best match of test cases
-        for broker_version, struct in sorted(test_cases, reverse=True):
-            if struct.API_KEY not in api_versions:
+        for broker_version, proto_struct in sorted(test_cases, reverse=True):
+            if proto_struct.API_KEY not in api_versions:
                 continue
-            min_version, max_version = api_versions[struct.API_KEY]
-            if min_version <= struct.API_VERSION <= max_version:
+            min_version, max_version = api_versions[proto_struct.API_KEY]
+            if min_version <= proto_struct.API_VERSION <= max_version:
                 return broker_version
 
         # We know that ApiVersionsResponse is only supported in 0.10+
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
index eefac5b..a833a5b 100644
--- a/kafka/consumer/fetcher.py
+++ b/kafka/consumer/fetcher.py
@@ -501,7 +501,7 @@ class Fetcher(six.Iterator):
         # If unpacking raises StopIteration, it is erroneously
         # caught by the generator. We want all exceptions to be raised
         # back to the user. See Issue 545
-        except StopIteration as e:
+        except StopIteration:
             log.exception('StopIteration raised unpacking messageset')
             raise RuntimeError('StopIteration raised unpacking messageset')
 
@@ -1001,7 +1001,7 @@ class FetchSessionHandler(object):
             log.debug("Built full fetch %s for node %s with %s partition(s).",
                 self.next_metadata, self.node_id, len(next_partitions))
             self.session_partitions = next_partitions
-            return FetchRequestData(next_partitions, None, self.next_metadata);
+            return FetchRequestData(next_partitions, None, self.next_metadata)
 
         prev_tps = set(self.session_partitions.keys())
         next_tps = set(next_partitions.keys())
diff --git a/kafka/coordinator/assignors/sticky/sticky_assignor.py b/kafka/coordinator/assignors/sticky/sticky_assignor.py
index dce714f..6e79c59 100644
--- a/kafka/coordinator/assignors/sticky/sticky_assignor.py
+++ b/kafka/coordinator/assignors/sticky/sticky_assignor.py
@@ -2,7 +2,6 @@ import logging
 from collections import defaultdict, namedtuple
 from copy import deepcopy
 
-from kafka.cluster import ClusterMetadata
 from kafka.coordinator.assignors.abstract import AbstractPartitionAssignor
 from kafka.coordinator.assignors.sticky.partition_movements import PartitionMovements
 from kafka.coordinator.assignors.sticky.sorted_set import SortedSet
diff --git a/kafka/coordinator/consumer.py b/kafka/coordinator/consumer.py
index 36c91ee..9c662ce 100644
--- a/kafka/coordinator/consumer.py
+++ b/kafka/coordinator/consumer.py
@@ -796,7 +796,7 @@ class ConsumerCoordinator(BaseCoordinator):
             elif error_type is Errors.GroupAuthorizationFailedError:
                 future.failure(error)
             else:
-                log.error("Unknown error fetching offsets for %s: %s", tp, error)
+                log.error("Unknown error fetching offsets: %s", error)
                 future.failure(error)
             return
 
diff --git a/kafka/metrics/metric_name.py b/kafka/metrics/metric_name.py
index b5acd16..32a7e3a 100644
--- a/kafka/metrics/metric_name.py
+++ b/kafka/metrics/metric_name.py
@@ -93,7 +93,7 @@ class MetricName(object):
             return True
         if other is None:
             return False
-        return (type(self) == type(other) and
+        return (isinstance(self, type(other)) and
                 self.group == other.group and
                 self.name == other.name and
                 self.tags == other.tags)
diff --git a/kafka/metrics/quota.py b/kafka/metrics/quota.py
index 4d1b0d6..237edf8 100644
--- a/kafka/metrics/quota.py
+++ b/kafka/metrics/quota.py
@@ -34,7 +34,7 @@ class Quota(object):
     def __eq__(self, other):
         if self is other:
             return True
-        return (type(self) == type(other) and
+        return (isinstance(self, type(other)) and
                 self.bound == other.bound and
                 self.is_upper_bound() == other.is_upper_bound())
 
diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index 668387a..2a70700 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -612,8 +612,8 @@ class KafkaProducer(object):
 
             if headers is None:
                 headers = []
-            assert type(headers) == list
-            assert all(type(item) == tuple and len(item) == 2 and type(item[0]) == str and type(item[1]) == bytes for item in headers)
+            assert isinstance(headers, list)
+            assert all(isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) and isinstance(item[1], bytes) for item in headers)
 
             message_size = self._estimate_size_in_bytes(key_bytes, value_bytes, headers)
             self._ensure_valid_record_size(message_size)
diff --git a/kafka/protocol/api_versions.py b/kafka/protocol/api_versions.py
index dc0aa58..7e2e612 100644
--- a/kafka/protocol/api_versions.py
+++ b/kafka/protocol/api_versions.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import
 
+from io import BytesIO
+
 from kafka.protocol.api import Request, Response
 from kafka.protocol.types import Array, Int16, Int32, Schema
 
diff --git a/kafka/protocol/commit.py b/kafka/protocol/commit.py
index 53c2466..a0439e7 100644
--- a/kafka/protocol/commit.py
+++ b/kafka/protocol/commit.py
@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.protocol.api import Request, Response
-from kafka.protocol.types import Array, Int8, Int16, Int32, Int64, Schema, String
+from kafka.protocol.types import Array, Int16, Int32, Int64, Schema, String
 
 
 class OffsetCommitResponse_v0(Response):
diff --git a/kafka/protocol/find_coordinator.py b/kafka/protocol/find_coordinator.py
index a68a239..be5b45d 100644
--- a/kafka/protocol/find_coordinator.py
+++ b/kafka/protocol/find_coordinator.py
@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.protocol.api import Request, Response
-from kafka.protocol.types import Array, Int8, Int16, Int32, Int64, Schema, String
+from kafka.protocol.types import Int8, Int16, Int32, Schema, String
 
 
 class FindCoordinatorResponse_v0(Response):
diff --git a/kafka/protocol/offset_for_leader_epoch.py b/kafka/protocol/offset_for_leader_epoch.py
index afe8284..8465588 100644
--- a/kafka/protocol/offset_for_leader_epoch.py
+++ b/kafka/protocol/offset_for_leader_epoch.py
@@ -4,7 +4,7 @@ from kafka.protocol.api import Request, Response
 from kafka.protocol.types import Array, CompactArray, CompactString, Int16, Int32, Int64, Schema, String, TaggedFields
 
 
-class OffsetForLeaderEpochResponse_v0(Request):
+class OffsetForLeaderEpochResponse_v0(Response):
     API_KEY = 23
     API_VERSION = 0
     SCHEMA = Schema(
@@ -16,7 +16,7 @@ class OffsetForLeaderEpochResponse_v0(Request):
                 ('end_offset', Int64))))))
 
 
-class OffsetForLeaderEpochResponse_v1(Request):
+class OffsetForLeaderEpochResponse_v1(Response):
     API_KEY = 23
     API_VERSION = 1
     SCHEMA = Schema(
@@ -29,7 +29,7 @@ class OffsetForLeaderEpochResponse_v1(Request):
                 ('end_offset', Int64))))))
 
 
-class OffsetForLeaderEpochResponse_v2(Request):
+class OffsetForLeaderEpochResponse_v2(Response):
     API_KEY = 23
     API_VERSION = 2
     SCHEMA = Schema(
@@ -43,13 +43,13 @@ class OffsetForLeaderEpochResponse_v2(Request):
                 ('end_offset', Int64))))))
 
 
-class OffsetForLeaderEpochResponse_v3(Request):
+class OffsetForLeaderEpochResponse_v3(Response):
     API_KEY = 23
     API_VERSION = 3
     SCHEMA = OffsetForLeaderEpochResponse_v2.SCHEMA
 
 
-class OffsetForLeaderEpochResponse_v4(Request):
+class OffsetForLeaderEpochResponse_v4(Response):
     API_KEY = 23
     API_VERSION = 4
     SCHEMA = Schema(
diff --git a/kafka/protocol/sasl_authenticate.py b/kafka/protocol/sasl_authenticate.py
index 528bb3c..a2b9b19 100644
--- a/kafka/protocol/sasl_authenticate.py
+++ b/kafka/protocol/sasl_authenticate.py
@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.protocol.api import Request, Response
-from kafka.protocol.types import Array, Bytes, Int16, Int64, Schema, String
+from kafka.protocol.types import Bytes, Int16, Int64, Schema, String
 
 
 class SaslAuthenticateResponse_v0(Response):
diff --git a/kafka/sasl/__init__.py b/kafka/sasl/__init__.py
index e36d1df..8677f60 100644
--- a/kafka/sasl/__init__.py
+++ b/kafka/sasl/__init__.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.sasl.gssapi import SaslMechanismGSSAPI
+from kafka.sasl.msk import SaslMechanismAwsMskIam
 from kafka.sasl.oauth import SaslMechanismOAuth
 from kafka.sasl.plain import SaslMechanismPlain
 from kafka.sasl.scram import SaslMechanismScram
@@ -24,3 +25,4 @@ register_sasl_mechanism('OAUTHBEARER', SaslMechanismOAuth)
 register_sasl_mechanism('PLAIN', SaslMechanismPlain)
 register_sasl_mechanism('SCRAM-SHA-256', SaslMechanismScram)
 register_sasl_mechanism('SCRAM-SHA-512', SaslMechanismScram)
+register_sasl_mechanism('AWS_MSK_IAM', SaslMechanismAwsMskIam)
diff --git a/kafka/sasl/gssapi.py b/kafka/sasl/gssapi.py
index 1be3de4..60b658c 100644
--- a/kafka/sasl/gssapi.py
+++ b/kafka/sasl/gssapi.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import
 
+import struct
+
 # needed for SASL_GSSAPI authentication:
 try:
     import gssapi
@@ -66,7 +68,7 @@ class SaslMechanismGSSAPI(SaslMechanism):
             # simply set QoP to 'auth' only (first octet). We reuse the max message size proposed
             # by the server
             message_parts = [
-                Int8.encode(self.SASL_QOP_AUTH & Int8.decode(io.BytesIO(msg[0:1]))),
+                struct.pack('>b', self.SASL_QOP_AUTH & struct.unpack('>b', msg[0:1])),
                 msg[:1],
                 self.auth_id.encode(),
             ]
diff --git a/kafka/sasl/msk.py b/kafka/sasl/msk.py
new file mode 100644
index 0000000..db56b48
--- /dev/null
+++ b/kafka/sasl/msk.py
@@ -0,0 +1,233 @@
+from __future__ import absolute_import
+
+import datetime
+import hashlib
+import hmac
+import json
+import string
+
+# needed for AWS_MSK_IAM authentication:
+try:
+    from botocore.session import Session as BotoSession
+except ImportError:
+    # no botocore available, will disable AWS_MSK_IAM mechanism
+    BotoSession = None
+
+from kafka.sasl.abc import SaslMechanism
+from kafka.vendor.six.moves import urllib
+
+
+class SaslMechanismAwsMskIam(SaslMechanism):
+    def __init__(self, **config):
+        assert BotoSession is not None, 'AWS_MSK_IAM requires the "botocore" package'
+        assert config.get('security_protocol', '') == 'SASL_SSL', 'AWS_MSK_IAM requires SASL_SSL'
+        assert 'host' in config, 'AWS_MSK_IAM requires host configuration'
+        self.host = config['host']
+        self._auth = None
+        self._is_done = False
+        self._is_authenticated = False
+
+    def auth_bytes(self):
+        session = BotoSession()
+        credentials = session.get_credentials().get_frozen_credentials()
+        client = AwsMskIamClient(
+            host=self.host,
+            access_key=credentials.access_key,
+            secret_key=credentials.secret_key,
+            region=session.get_config_variable('region'),
+            token=credentials.token,
+        )
+        return client.first_message()
+
+    def receive(self, auth_bytes):
+        self._is_done = True
+        self._is_authenticated = auth_bytes != b''
+        self._auth = auth_bytes.deode('utf-8')
+
+    def is_done(self):
+        return self._is_done
+
+    def is_authenticated(self):
+        return self._is_authenticated
+
+    def auth_details(self):
+        if not self.is_authenticated:
+            raise RuntimeError('Not authenticated yet!')
+        return 'Authenticated via SASL / AWS_MSK_IAM %s' % (self._auth,)
+
+
+class AwsMskIamClient:
+    UNRESERVED_CHARS = string.ascii_letters + string.digits + '-._~'
+
+    def __init__(self, host, access_key, secret_key, region, token=None):
+        """
+        Arguments:
+            host (str): The hostname of the broker.
+            access_key (str): An AWS_ACCESS_KEY_ID.
+            secret_key (str): An AWS_SECRET_ACCESS_KEY.
+            region (str): An AWS_REGION.
+            token (Optional[str]): An AWS_SESSION_TOKEN if using temporary
+                credentials.
+        """
+        self.algorithm = 'AWS4-HMAC-SHA256'
+        self.expires = '900'
+        self.hashfunc = hashlib.sha256
+        self.headers = [
+            ('host', host)
+        ]
+        self.version = '2020_10_22'
+
+        self.service = 'kafka-cluster'
+        self.action = '{}:Connect'.format(self.service)
+
+        now = datetime.datetime.utcnow()
+        self.datestamp = now.strftime('%Y%m%d')
+        self.timestamp = now.strftime('%Y%m%dT%H%M%SZ')
+
+        self.host = host
+        self.access_key = access_key
+        self.secret_key = secret_key
+        self.region = region
+        self.token = token
+
+    @property
+    def _credential(self):
+        return '{0.access_key}/{0._scope}'.format(self)
+
+    @property
+    def _scope(self):
+        return '{0.datestamp}/{0.region}/{0.service}/aws4_request'.format(self)
+
+    @property
+    def _signed_headers(self):
+        """
+        Returns (str):
+            An alphabetically sorted, semicolon-delimited list of lowercase
+            request header names.
+        """
+        return ';'.join(sorted(k.lower() for k, _ in self.headers))
+
+    @property
+    def _canonical_headers(self):
+        """
+        Returns (str):
+            A newline-delited list of header names and values.
+            Header names are lowercased.
+        """
+        return '\n'.join(map(':'.join, self.headers)) + '\n'
+
+    @property
+    def _canonical_request(self):
+        """
+        Returns (str):
+            An AWS Signature Version 4 canonical request in the format:
+                <Method>\n
+                <Path>\n
+                <CanonicalQueryString>\n
+                <CanonicalHeaders>\n
+                <SignedHeaders>\n
+                <HashedPayload>
+        """
+        # The hashed_payload is always an empty string for MSK.
+        hashed_payload = self.hashfunc(b'').hexdigest()
+        return '\n'.join((
+            'GET',
+            '/',
+            self._canonical_querystring,
+            self._canonical_headers,
+            self._signed_headers,
+            hashed_payload,
+        ))
+
+    @property
+    def _canonical_querystring(self):
+        """
+        Returns (str):
+            A '&'-separated list of URI-encoded key/value pairs.
+        """
+        params = []
+        params.append(('Action', self.action))
+        params.append(('X-Amz-Algorithm', self.algorithm))
+        params.append(('X-Amz-Credential', self._credential))
+        params.append(('X-Amz-Date', self.timestamp))
+        params.append(('X-Amz-Expires', self.expires))
+        if self.token:
+            params.append(('X-Amz-Security-Token', self.token))
+        params.append(('X-Amz-SignedHeaders', self._signed_headers))
+
+        return '&'.join(self._uriencode(k) + '=' + self._uriencode(v) for k, v in params)
+
+    @property
+    def _signing_key(self):
+        """
+        Returns (bytes):
+            An AWS Signature V4 signing key generated from the secret_key, date,
+            region, service, and request type.
+        """
+        key = self._hmac(('AWS4' + self.secret_key).encode('utf-8'), self.datestamp)
+        key = self._hmac(key, self.region)
+        key = self._hmac(key, self.service)
+        key = self._hmac(key, 'aws4_request')
+        return key
+
+    @property
+    def _signing_str(self):
+        """
+        Returns (str):
+            A string used to sign the AWS Signature V4 payload in the format:
+                <Algorithm>\n
+                <Timestamp>\n
+                <Scope>\n
+                <CanonicalRequestHash>
+        """
+        canonical_request_hash = self.hashfunc(self._canonical_request.encode('utf-8')).hexdigest()
+        return '\n'.join((self.algorithm, self.timestamp, self._scope, canonical_request_hash))
+
+    def _uriencode(self, msg):
+        """
+        Arguments:
+            msg (str): A string to URI-encode.
+
+        Returns (str):
+            The URI-encoded version of the provided msg, following the encoding
+            rules specified: https://github.com/aws/aws-msk-iam-auth#uriencode
+        """
+        return urllib.parse.quote(msg, safe=self.UNRESERVED_CHARS)
+
+    def _hmac(self, key, msg):
+        """
+        Arguments:
+            key (bytes): A key to use for the HMAC digest.
+            msg (str): A value to include in the HMAC digest.
+        Returns (bytes):
+            An HMAC digest of the given key and msg.
+        """
+        return hmac.new(key, msg.encode('utf-8'), digestmod=self.hashfunc).digest()
+
+    def first_message(self):
+        """
+        Returns (bytes):
+            An encoded JSON authentication payload that can be sent to the
+            broker.
+        """
+        signature = hmac.new(
+            self._signing_key,
+            self._signing_str.encode('utf-8'),
+            digestmod=self.hashfunc,
+        ).hexdigest()
+        msg = {
+            'version':  self.version,
+            'host': self.host,
+            'user-agent': 'kafka-python',
+            'action': self.action,
+            'x-amz-algorithm': self.algorithm,
+            'x-amz-credential': self._credential,
+            'x-amz-date': self.timestamp,
+            'x-amz-signedheaders': self._signed_headers,
+            'x-amz-expires': self.expires,
+            'x-amz-signature': signature,
+        }
+        if self.token:
+            msg['x-amz-security-token'] = self.token
+
+        return json.dumps(msg, separators=(',', ':')).encode('utf-8')
diff --git a/pyproject.toml b/pyproject.toml
index 48be87f..ddd40a0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,7 @@ crc32c = ["crc32c"]
 lz4 = ["lz4"]
 snappy = ["python-snappy"]
 zstd = ["zstandard"]
-testing = ["pytest", "mock", "pytest-mock"]
+testing = ["pytest", "mock; python_version < '3.3'", "pytest-mock"]
 
 [tool.setuptools]
 include-package-data = false
diff --git a/requirements-dev.txt b/requirements-dev.txt
index e49608a..6cfb6d8 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -3,7 +3,7 @@ crc32c
 docker-py
 flake8
 lz4
-mock
+mock; python_version < '3.3'
 py
 pylint
 pytest


Test patch:
diff --git a/test/record/test_default_records.py b/test/record/test_default_records.py
index c3a7b02..e1c840f 100644
--- a/test/record/test_default_records.py
+++ b/test/record/test_default_records.py
@@ -1,7 +1,10 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 import pytest
-from mock import patch
+try:
+    from unittest.mock import patch
+except ImportError:
+    from mock import patch
 import kafka.codec
 from kafka.record.default_records import (
     DefaultRecordBatch, DefaultRecordBatchBuilder
diff --git a/test/record/test_legacy_records.py b/test/record/test_legacy_records.py
index 43970f7..b15b537 100644
--- a/test/record/test_legacy_records.py
+++ b/test/record/test_legacy_records.py
@@ -1,6 +1,9 @@
 from __future__ import unicode_literals
 import pytest
-from mock import patch
+try:
+    from unittest.mock import patch
+except ImportError:
+    from mock import patch
 from kafka.record.legacy_records import (
     LegacyRecordBatch, LegacyRecordBatchBuilder
 )
diff --git a/test/sasl/test_msk.py b/test/sasl/test_msk.py
new file mode 100644
index 0000000..297ca84
--- /dev/null
+++ b/test/sasl/test_msk.py
@@ -0,0 +1,67 @@
+import datetime
+import json
+
+from kafka.sasl.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.sasl.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
diff --git a/test/test_admin.py b/test/test_admin.py
index 279f85a..cdb7424 100644
--- a/test/test_admin.py
+++ b/test/test_admin.py
@@ -6,7 +6,7 @@ from kafka.errors import IllegalArgumentError
 
 def test_config_resource():
     with pytest.raises(KeyError):
-        bad_resource = kafka.admin.ConfigResource('something', 'foo')
+        _bad_resource = kafka.admin.ConfigResource('something', 'foo')
     good_resource = kafka.admin.ConfigResource('broker', 'bar')
     assert good_resource.resource_type == kafka.admin.ConfigResourceType.BROKER
     assert good_resource.name == 'bar'
@@ -59,11 +59,11 @@ def test_acl_resource():
 
 def test_new_topic():
     with pytest.raises(IllegalArgumentError):
-        bad_topic = kafka.admin.NewTopic('foo', -1, -1)
+        _bad_topic = kafka.admin.NewTopic('foo', -1, -1)
     with pytest.raises(IllegalArgumentError):
-        bad_topic = kafka.admin.NewTopic('foo', 1, -1)
+        _bad_topic = kafka.admin.NewTopic('foo', 1, -1)
     with pytest.raises(IllegalArgumentError):
-        bad_topic = kafka.admin.NewTopic('foo', 1, 1, {1: [1, 1, 1]})
+        _bad_topic = kafka.admin.NewTopic('foo', 1, 1, {1: [1, 1, 1]})
     good_topic = kafka.admin.NewTopic('foo', 1, 2)
     assert good_topic.name == 'foo'
     assert good_topic.num_partitions == 1
diff --git a/test/test_admin_integration.py b/test/test_admin_integration.py
index bd2fd21..2f6b765 100644
--- a/test/test_admin_integration.py
+++ b/test/test_admin_integration.py
@@ -140,7 +140,7 @@ def test_describe_configs_invalid_broker_id_raises(kafka_admin_client):
     broker_id = "str"
 
     with pytest.raises(ValueError):
-        configs = kafka_admin_client.describe_configs([ConfigResource(ConfigResourceType.BROKER, broker_id)])
+        kafka_admin_client.describe_configs([ConfigResource(ConfigResourceType.BROKER, broker_id)])
 
 
 @pytest.mark.skipif(env_kafka_version() < (0, 11), reason='Describe consumer group requires broker >=0.11')
@@ -148,7 +148,7 @@ def test_describe_consumer_group_does_not_exist(kafka_admin_client):
     """Tests that the describe consumer group call fails if the group coordinator is not available
     """
     with pytest.raises(GroupCoordinatorNotAvailableError):
-        group_description = kafka_admin_client.describe_consumer_groups(['test'])
+        kafka_admin_client.describe_consumer_groups(['test'])
 
 
 @pytest.mark.skipif(env_kafka_version() < (0, 11), reason='Describe consumer group requires broker >=0.11')
diff --git a/test/test_cluster.py b/test/test_cluster.py
index b55bdc5..f0a2f83 100644
--- a/test/test_cluster.py
+++ b/test/test_cluster.py
@@ -1,8 +1,6 @@
 # pylint: skip-file
 from __future__ import absolute_import
 
-import pytest
-
 from kafka.cluster import ClusterMetadata
 from kafka.protocol.metadata import MetadataResponse
 
diff --git a/test/test_conn.py b/test/test_conn.py
index 959cbb4..ea88fd0 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -4,7 +4,10 @@ from __future__ import absolute_import
 from errno import EALREADY, EINPROGRESS, EISCONN, ECONNRESET
 import socket
 
-import mock
+try:
+    from unittest import mock
+except ImportError:
+    import mock
 import pytest
 
 from kafka.conn import BrokerConnection, ConnectionStates, collect_hosts
diff --git a/test/test_consumer_integration.py b/test/test_consumer_integration.py
index 5aeb63d..af8ec68 100644
--- a/test/test_consumer_integration.py
+++ b/test/test_consumer_integration.py
@@ -1,7 +1,10 @@
 import logging
 import time
 
-from mock import patch, ANY
+try:
+    from unittest.mock import patch, ANY
+except ImportError:
+    from mock import patch, ANY
 import pytest
 from kafka.vendor.six.moves import range
 
diff --git a/test/test_fetcher.py b/test/test_fetcher.py
index 7e948e3..479f6e2 100644
--- a/test/test_fetcher.py
+++ b/test/test_fetcher.py
@@ -10,7 +10,7 @@ import time
 
 from kafka.client_async import KafkaClient
 from kafka.consumer.fetcher import (
-    CompletedFetch, ConsumerRecord, Fetcher, NoOffsetForPartitionError
+    CompletedFetch, ConsumerRecord, Fetcher
 )
 from kafka.consumer.subscription_state import SubscriptionState
 import kafka.errors as Errors
diff --git a/test/test_object_conversion.py b/test/test_object_conversion.py
index 9b1ff21..a48eb06 100644
--- a/test/test_object_conversion.py
+++ b/test/test_object_conversion.py
@@ -207,7 +207,7 @@ def test_with_metadata_response():
     assert len(obj['topics']) == 2
     assert obj['topics'][0]['error_code'] == 0
     assert obj['topics'][0]['topic'] == 'testtopic1'
-    assert obj['topics'][0]['is_internal'] == False
+    assert obj['topics'][0]['is_internal'] is False
     assert len(obj['topics'][0]['partitions']) == 2
     assert obj['topics'][0]['partitions'][0]['error_code'] == 0
     assert obj['topics'][0]['partitions'][0]['partition'] == 0
@@ -224,7 +224,7 @@ def test_with_metadata_response():
 
     assert obj['topics'][1]['error_code'] == 0
     assert obj['topics'][1]['topic'] == 'other-test-topic'
-    assert obj['topics'][1]['is_internal'] == True
+    assert obj['topics'][1]['is_internal'] is True
     assert len(obj['topics'][1]['partitions']) == 1
     assert obj['topics'][1]['partitions'][0]['error_code'] == 0
     assert obj['topics'][1]['partitions'][0]['partition'] == 0
diff --git a/test/test_protocol.py b/test/test_protocol.py
index 6f94c74..d0cc7ed 100644
--- a/test/test_protocol.py
+++ b/test/test_protocol.py
@@ -2,8 +2,6 @@
 import io
 import struct
 
-import pytest
-
 from kafka.protocol.api import RequestHeader
 from kafka.protocol.fetch import FetchRequest, FetchResponse
 from kafka.protocol.find_coordinator import FindCoordinatorRequest
@@ -273,7 +271,7 @@ def test_decode_fetch_response_partial():
 
 def test_struct_unrecognized_kwargs():
     try:
-        mr = MetadataRequest[0](topicz='foo')
+        _mr = MetadataRequest[0](topicz='foo')
         assert False, 'Structs should not allow unrecognized kwargs'
     except ValueError:
         pass
@@ -331,6 +329,6 @@ def test_compact_data_structs():
     assert CompactBytes.decode(io.BytesIO(b'\x00')) is None
     enc = CompactBytes.encode(b'')
     assert enc == b'\x01'
-    assert CompactBytes.decode(io.BytesIO(b'\x01')) is b''
+    assert CompactBytes.decode(io.BytesIO(b'\x01')) == b''
     enc = CompactBytes.encode(b'foo')
     assert CompactBytes.decode(io.BytesIO(enc)) == b'foo'
diff --git a/test/test_sender.py b/test/test_sender.py
index 83a26cd..3da1a9f 100644
--- a/test/test_sender.py
+++ b/test/test_sender.py
@@ -5,7 +5,6 @@ import pytest
 import io
 
 from kafka.client_async import KafkaClient
-from kafka.cluster import ClusterMetadata
 from kafka.metrics import Metrics
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
 from kafka.protocol.produce import ProduceRequest