diff --git a/providers/standard/src/airflow/providers/standard/utils/python_virtualenv.py b/providers/standard/src/airflow/providers/standard/utils/python_virtualenv.py index 891b3e0bce3a1..33b4a0fb7e0fc 100644 --- a/providers/standard/src/airflow/providers/standard/utils/python_virtualenv.py +++ b/providers/standard/src/airflow/providers/standard/utils/python_virtualenv.py @@ -200,9 +200,10 @@ def prepare_virtualenv( if _use_uv(): venv_cmd = _generate_uv_cmd(venv_directory, python_bin, system_site_packages) + _execute_in_subprocess(venv_cmd, env={**os.environ, **_index_urls_to_uv_env_vars(index_urls)}) else: venv_cmd = _generate_venv_cmd(venv_directory, python_bin, system_site_packages) - _execute_in_subprocess(venv_cmd) + _execute_in_subprocess(venv_cmd) pip_cmd = None if requirements is not None and len(requirements) != 0: diff --git a/providers/standard/tests/unit/standard/operators/test_python.py b/providers/standard/tests/unit/standard/operators/test_python.py index cbc2bca2958a0..bccd232479308 100644 --- a/providers/standard/tests/unit/standard/operators/test_python.py +++ b/providers/standard/tests/unit/standard/operators/test_python.py @@ -1420,9 +1420,18 @@ def f(a): @mock.patch( "airflow.providers.standard.utils.python_virtualenv._execute_in_subprocess", - wraps=_execute_in_subprocess, ) - def test_with_index_urls(self, wrapped_execute_in_subprocess): + def test_with_index_urls(self, mock_execute_in_subprocess): + def safe_execute_in_subprocess(cmd, **kwargs): + """Wrapper that removes unreachable URLs from env before executing.""" + env = kwargs.get("env", {}).copy() + # Remove fake URLs to allow venv creation to succeed + env.pop("UV_DEFAULT_INDEX", None) + env.pop("UV_INDEX", None) + return _execute_in_subprocess(cmd, **{**kwargs, "env": env if env else None}) + + mock_execute_in_subprocess.side_effect = safe_execute_in_subprocess + def f(a): import sys from pathlib import Path @@ -1443,15 +1452,30 @@ def f(a): op_args=[4], ) - # first call creates venv, second call installs packages - package_install_call_args = wrapped_execute_in_subprocess.call_args[1] - assert package_install_call_args["env"]["UV_DEFAULT_INDEX"] == "https://first.package.index" + # Verify that env was passed with the correct index URLs to pip install + # The first call is venv creation (with cleaned env), second call is pip install (with full env) + assert len(mock_execute_in_subprocess.call_args_list) >= 2 + package_install_call_args = mock_execute_in_subprocess.call_args_list[1] + assert package_install_call_args[1]["env"]["UV_DEFAULT_INDEX"] == "https://first.package.index" assert ( - package_install_call_args["env"]["UV_INDEX"] + package_install_call_args[1]["env"]["UV_INDEX"] == "http://second.package.index http://third.package.index" ) - def test_with_index_url_from_connection(self, monkeypatch): + @mock.patch( + "airflow.providers.standard.utils.python_virtualenv._execute_in_subprocess", + ) + def test_with_index_url_from_connection(self, mock_execute_in_subprocess, monkeypatch): + def safe_execute_in_subprocess(cmd, **kwargs): + """Wrapper that removes unreachable URLs from env before executing.""" + env = kwargs.get("env", {}).copy() + # Remove fake URLs to allow venv creation to succeed + env.pop("UV_DEFAULT_INDEX", None) + env.pop("UV_INDEX", None) + return _execute_in_subprocess(cmd, **{**kwargs, "env": env if env else None}) + + mock_execute_in_subprocess.side_effect = safe_execute_in_subprocess + class MockConnection(Connection): """Mock for the Connection class.""" diff --git a/providers/standard/tests/unit/standard/utils/test_python_virtualenv.py b/providers/standard/tests/unit/standard/utils/test_python_virtualenv.py index 2aaa9a1283d3f..e35ee2eea9958 100644 --- a/providers/standard/tests/unit/standard/utils/test_python_virtualenv.py +++ b/providers/standard/tests/unit/standard/utils/test_python_virtualenv.py @@ -94,7 +94,8 @@ def test_should_create_virtualenv_uv(self, mock_execute_in_subprocess): ) assert python_bin == "/VENV/bin/python" mock_execute_in_subprocess.assert_called_once_with( - ["uv", "venv", "--allow-existing", "--seed", "--python", "pythonVER", "/VENV"] + ["uv", "venv", "--allow-existing", "--seed", "--python", "pythonVER", "/VENV"], + env=mock.ANY, ) @mock.patch("airflow.providers.standard.utils.python_virtualenv._execute_in_subprocess") @@ -125,7 +126,8 @@ def test_should_create_virtualenv_with_system_packages_uv(self, mock_execute_in_ "pythonVER", "--system-site-packages", "/VENV", - ] + ], + env=mock.ANY, ) @mock.patch("airflow.providers.standard.utils.python_virtualenv._execute_in_subprocess") @@ -205,6 +207,48 @@ def test_should_create_virtualenv_with_extra_packages_uv(self, mock_execute_in_s env=mock.ANY, ) + @mock.patch("airflow.providers.standard.utils.python_virtualenv._generate_pip_conf") + @mock.patch("airflow.providers.standard.utils.python_virtualenv._execute_in_subprocess") + @conf_vars({("standard", "venv_install_method"): "uv"}) + def test_venv_creation_with_index_urls( + self, mock_execute_in_subprocess, mock_generate_pip_conf, tmp_path: Path + ): + """ + Test that uv venv creation passes UV_DEFAULT_INDEX env var. + + When package-index connections are available, UV_DEFAULT_INDEX and UV_INDEX + environment variables are passed to uv venv command to ensure packages can be + downloaded from the specified index during venv creation (e.g., when installing + seed packages like pip, setuptools, wheel). + """ + venv_dir = str(tmp_path / "venv") + python_bin = prepare_virtualenv( + venv_directory=venv_dir, + python_bin="pythonVER", + system_site_packages=False, + requirements=["somepackage"], + index_urls=["https://private.package.index"], + ) + assert python_bin == f"{venv_dir}/bin/python" + + # First call: venv creation should have UV_DEFAULT_INDEX in env + venv_call_args = mock_execute_in_subprocess.call_args_list[0] + venv_cmd = venv_call_args[0][0] + venv_kwargs = venv_call_args[1] + + assert venv_cmd == ["uv", "venv", "--allow-existing", "--seed", "--python", "pythonVER", venv_dir] + assert "env" in venv_kwargs + assert venv_kwargs["env"]["UV_DEFAULT_INDEX"] == "https://private.package.index" + + # Second call: pip install should also have UV_DEFAULT_INDEX + pip_call_args = mock_execute_in_subprocess.call_args_list[1] + pip_cmd = pip_call_args[0][0] + pip_kwargs = pip_call_args[1] + + assert pip_cmd == ["uv", "pip", "install", "--python", f"{venv_dir}/bin/python", "somepackage"] + assert "env" in pip_kwargs + assert pip_kwargs["env"]["UV_DEFAULT_INDEX"] == "https://private.package.index" + @pytest.mark.parametrize( ("decorators", "expected_decorators"), [