From c662d941b13f11620f60fc48580f2f6643051ca6 Mon Sep 17 00:00:00 2001
From: Jonas Keeling <jonas.keeling@aiven.io>
Date: Mon, 26 Feb 2024 14:36:42 +0100
Subject: [PATCH] bugfix: catch KeyError in check_version and retry

---
 kafka/client_async.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/kafka/client_async.py b/kafka/client_async.py
index 2e3a680ff..9a720112d 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -912,7 +912,13 @@ def check_version(self, node_id=None, timeout=2, strict=False):
                 self._lock.release()
                 raise Errors.NoBrokersAvailable()
             self._maybe_connect(try_node)
-            conn = self._conns[try_node]
+            try:
+                conn = self._conns[try_node]
+            except KeyError:
+                if node_id is not None:
+                    self._lock.release()
+                    raise Errors.NodeNotReadyError()
+                continue
 
             # We will intentionally cause socket failures
             # These should not trigger metadata refresh