diff --git a/tests/common.py b/tests/common.py
index eeec1acab..c24cb2435 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -387,9 +387,18 @@ def zigpy_device_from_device_data(
 
         for cluster_type in ("in_clusters", "out_clusters"):
             for cluster in ep[cluster_type]:
-                real_cluster = getattr(endpoint, cluster_type)[
+                real_cluster = getattr(endpoint, cluster_type).get(
                     int(cluster["cluster_id"], 16)
-                ]
+                )
+
+                if real_cluster is None:
+                    _LOGGER.warning(
+                        "Cluster %s not found in endpoint %s of device %s",
+                        cluster["cluster_id"],
+                        epid,
+                        device.ieee,
+                    )
+                    continue
 
                 if patch_cluster:
                     patch_cluster_for_testing(real_cluster)
diff --git a/tools/import_legacy_diagnostics.py b/tools/import_legacy_diagnostics.py
index 7c7b14fdc..5ea774e0c 100644
--- a/tools/import_legacy_diagnostics.py
+++ b/tools/import_legacy_diagnostics.py
@@ -105,6 +105,7 @@ def zigpy_device_from_legacy_diagnostics(
     for epid, ep in cluster_data.items():
         endpoint.request = AsyncMock(return_value=[0])
         for cluster_id, cluster in ep["in_clusters"].items():
+            # todo: likely also add here?
             real_cluster = device.endpoints[int(epid)].in_clusters[int(cluster_id, 16)]
             if patch_cluster:
                 patch_cluster_for_testing(real_cluster)