From 8bd39cbbf2ddf649445c715248fd4de2b16637c9 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Sat, 9 Nov 2024 10:22:10 +0100
Subject: [PATCH] Fix remove_local_wheels for no build wheels (#798)

---
 builder/infra.py    |  2 +-
 tests/test_infra.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/builder/infra.py b/builder/infra.py
index c6914bf..7e425f3 100644
--- a/builder/infra.py
+++ b/builder/infra.py
@@ -162,6 +162,6 @@ def remove_local_wheels(
     for binary in exists:
         version = binary_package_map[binary]
         print(f"Found existing wheels for {binary}, removing local copy {version}")
-        for wheel in wheel_map[binary]:
+        for wheel in wheel_map.get(binary, ()):
             print(f"Removing local wheel {wheel}")
             wheel.unlink()
diff --git a/tests/test_infra.py b/tests/test_infra.py
index 2646958..9dc01e1 100644
--- a/tests/test_infra.py
+++ b/tests/test_infra.py
@@ -259,3 +259,18 @@ def test_remove_local_wheel_normalized_package_names(tmppath: Path) -> None:
 
     # grpcio and google-cloud-pubsub are removed
     assert {p.name for p in tmppath.glob("*.whl")} == set()
+
+
+def test_remove_local_wheel_no_build_wheels(tmppath: Path) -> None:
+    """Test remove_local_wheels does not fail with skip_exists and no build wheels."""
+    package_index = infra.extract_packages_from_index("https://example.com")
+    assert {p.name for p in tmppath.glob("*.whl")} == set()
+
+    infra.remove_local_wheels(
+        package_index,
+        skip_exists="grpcio",
+        packages=[
+            "grpcio==1.31.0",  # Exists in index
+        ],
+        wheels_dir=tmppath,
+    )