DO NOT MERGE: Add support for extra index url to the pip installer

homeassistant/requirements.py

@@ -88,7 +88,9 @@ def pip_kwargs(config_dir: str | None) -> dict[str, Any]:
         "no_cache_dir": is_docker,
         "timeout": PIP_TIMEOUT,
     }
-    if "WHEELS_LINKS" in os.environ:
+    if "EXTRA_INDEX_URL" in os.environ:
+        kwargs["extra_index_url"] = os.environ["EXTRA_INDEX_URL"]
+    elif "WHEELS_LINKS" in os.environ:
         kwargs["find_links"] = os.environ["WHEELS_LINKS"]
     if not (config_dir is None or pkg_util.is_virtual_env()) and not is_docker:
         kwargs["target"] = os.path.join(config_dir, "deps")

homeassistant/util/package.py

@@ -70,6 +70,7 @@ def install_package(
     find_links: str | None = None,
     timeout: int | None = None,
     no_cache_dir: bool | None = False,
+    extra_index_url: str | None = None,
 ) -> bool:
     """Install a package on PyPi. Accepts pip compatible package strings.
 
@@ -87,6 +88,8 @@ def install_package(
         args.append("--upgrade")
     if constraints is not None:
         args += ["--constraint", constraints]
+    if extra_index_url is not None:
+        args += ["--extra-index-url", extra_index_url, "--prefer-binary"]
     if find_links is not None:
         args += ["--find-links", find_links, "--prefer-binary"]
     if target:

tests/test_requirements.py

@@ -21,7 +21,8 @@
 def env_without_wheel_links():
     """Return env without wheel links."""
     env = dict(os.environ)
-    env.pop("WHEEL_LINKS", None)
+    env.pop("WHEELS_LINKS", None)
+    env.pop("EXTRA_INDEX_URL", None)
     return env
 
 
@@ -361,8 +362,8 @@ async def test_get_custom_integration_with_missing_after_dependencies(
     assert integration.domain == "test_custom_component"
 
 
-async def test_install_with_wheels_index(hass: HomeAssistant) -> None:
-    """Test an install attempt with wheels index URL."""
+async def test_install_with_wheels_links(hass: HomeAssistant) -> None:
+    """Test an install attempt with wheels links URL."""
     hass.config.skip_pip = False
     mock_integration(hass, MockModule("comp", requirements=["hello==1.0.0"]))
 
@@ -386,6 +387,31 @@ async def test_install_with_wheels_index(hass: HomeAssistant) -> None:
         )
 
 
+async def test_install_with_extra_index(hass: HomeAssistant) -> None:
+    """Test an install attempt with wheels index URL."""
+    hass.config.skip_pip = False
+    mock_integration(hass, MockModule("comp", requirements=["hello==1.0.0"]))
+
+    with patch("homeassistant.util.package.is_installed", return_value=False), patch(
+        "homeassistant.util.package.is_docker_env", return_value=True
+    ), patch("homeassistant.util.package.install_package") as mock_inst, patch.dict(
+        os.environ, {"EXTRA_INDEX_URL": "https://wheels.hass.io/test"}
+    ), patch(
+        "os.path.dirname"
+    ) as mock_dir:
+        mock_dir.return_value = "ha_package_path"
+        assert await setup.async_setup_component(hass, "comp", {})
+        assert "comp" in hass.config.components
+
+        assert mock_inst.call_args == call(
+            "hello==1.0.0",
+            extra_index_url="https://wheels.hass.io/test",
+            constraints=os.path.join("ha_package_path", CONSTRAINT_FILE),
+            timeout=60,
+            no_cache_dir=True,
+        )
+
+
 async def test_install_on_docker(hass: HomeAssistant) -> None:
     """Test an install attempt on an docker system env."""
     hass.config.skip_pip = False

tests/util/test_package.py

@@ -221,6 +221,35 @@ def test_install_find_links(mock_sys, mock_popen, mock_env_copy, mock_venv) -> N
     assert mock_popen.return_value.communicate.call_count == 1
 
 
+def test_install_extra_index_url(
+    mock_sys, mock_popen, mock_env_copy, mock_venv
+) -> None:
+    """Test install with extra_index_url on not installed package."""
+    env = mock_env_copy()
+    link = "https://wheels-repository"
+    assert package.install_package(TEST_NEW_REQ, False, extra_index_url=link)
+    assert mock_popen.call_count == 2
+    assert mock_popen.mock_calls[0] == call(
+        [
+            mock_sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "--quiet",
+            TEST_NEW_REQ,
+            "--extra-index-url",
+            link,
+            "--prefer-binary",
+        ],
+        stdin=PIPE,
+        stdout=PIPE,
+        stderr=PIPE,
+        env=env,
+        close_fds=False,
+    )
+    assert mock_popen.return_value.communicate.call_count == 1
+
+
 async def test_async_get_user_site(mock_env_copy) -> None:
     """Test async get user site directory."""
     deps_dir = "/deps_dir"