diff --git a/installer/__main__.py b/installer/__main__.py index b432fd9..fae2947 100644 --- a/installer/__main__.py +++ b/installer/__main__.py @@ -273,7 +273,7 @@ def bootstrap(): venv.import_package_bundle( server_site_packages_path, condition_env="DEEPNOTE_INCLUDE_SERVER_PACKAGES" ) - venv.import_package_bundle(system_site_packages_path) + venv.import_package_bundle(system_site_packages_path, priority=True) venv.import_package_bundle(kernel_site_package_path) # Phase 7.2: Symlink notebook to jupyter_server for compatibility with diff --git a/installer/module/virtual_environment.py b/installer/module/virtual_environment.py index d78799d..841409a 100644 --- a/installer/module/virtual_environment.py +++ b/installer/module/virtual_environment.py @@ -81,15 +81,29 @@ def start_server(self, command: str, cwd: Optional[str] = None) -> ServerProcess return server_proc def import_package_bundle( - self, bundle_site_package_path: str, condition_env: Optional[str] = None + self, + bundle_site_package_path: str, + condition_env: Optional[str] = None, + *, + priority: bool = False, ) -> None: """ Import a package bundle to the virtual environment. :param bundle_site_package_path: Absolute path to the package bundle. :param condition_env: Optional environment variable name. If provided, the bundle - will only be loaded when this env var is set to 'true'. + will only be loaded when this env var is set to 'true'. + Cannot be combined with priority. + :param priority: If True, insert at front of sys.path to override other bundles. + Cannot be combined with condition_env. + :raises ValueError: If both condition_env and priority are specified. """ + if condition_env and priority: + raise ValueError( + "condition_env and priority are mutually exclusive; " + "specify only one of them" + ) + pth_file_path = os.path.join(self._site_packages_path, "deepnote.pth") with open(pth_file_path, "a+", encoding="utf-8") as pth_file: @@ -105,6 +119,15 @@ def import_package_bundle( "Bundle was conditionally imported to the virtual environment " f"(loads when {condition_env}=true)." ) + elif priority: + # Insert at front of sys.path for higher priority (overrides other bundles) + pth_content = ( + f"import sys; sys.path.insert(0, '{bundle_site_package_path}')" + ) + pth_file.write(pth_content + "\n") + logger.info( + "Bundle was imported with priority to the virtual environment." + ) else: pth_file.write(bundle_site_package_path + "\n") logger.info( diff --git a/tests/unit/test_virtual_environment.py b/tests/unit/test_virtual_environment.py new file mode 100644 index 0000000..841d0d9 --- /dev/null +++ b/tests/unit/test_virtual_environment.py @@ -0,0 +1,100 @@ +"""Tests for VirtualEnvironment class.""" + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from installer.module.virtual_environment import VirtualEnvironment + + +class TestImportPackageBundle: + """Tests for import_package_bundle method.""" + + @pytest.fixture + def venv(self, tmp_path): + """Create a VirtualEnvironment with a temporary path.""" + venv_path = tmp_path / "venv" + with patch( + "installer.module.virtual_environment.get_current_python_version", + return_value="3.11", + ): + venv = VirtualEnvironment(str(venv_path)) + # Create the site-packages directory + Path(venv.site_packages_path).mkdir(parents=True, exist_ok=True) + return venv + + def test_import_package_bundle_plain_path(self, venv): + """Test that plain path is appended to .pth file.""" + bundle_path = "/some/bundle/site-packages" + + venv.import_package_bundle(bundle_path) + + pth_file = Path(venv.site_packages_path) / "deepnote.pth" + content = pth_file.read_text() + + assert content == f"{bundle_path}\n" + + def test_import_package_bundle_with_condition_env(self, venv): + """Test that conditional import uses sys.path.insert(0, ...).""" + bundle_path = "/server/libs/site-packages" + condition_env = "DEEPNOTE_INCLUDE_SERVER_PACKAGES" + + venv.import_package_bundle(bundle_path, condition_env=condition_env) + + pth_file = Path(venv.site_packages_path) / "deepnote.pth" + content = pth_file.read_text() + + assert "import os, sys" in content + assert f"sys.path.insert(0, '{bundle_path}')" in content + assert f"os.environ.get('{condition_env}', '').lower() == 'true'" in content + + def test_import_package_bundle_with_priority(self, venv): + """Test that priority=True uses sys.path.insert(0, ...).""" + bundle_path = "/usr/local/lib/python3.11/site-packages" + + venv.import_package_bundle(bundle_path, priority=True) + + pth_file = Path(venv.site_packages_path) / "deepnote.pth" + content = pth_file.read_text() + + assert "import sys" in content + assert f"sys.path.insert(0, '{bundle_path}')" in content + assert "os.environ.get" not in content + + def test_import_package_bundle_ordering(self, venv): + """Test .pth file content matches expected order from __main__.py.""" + server_libs = "/tmp/python3.11/server-libs/lib/python3.11/site-packages" + system_site = "/usr/local/lib/python3.11/site-packages" + kernel_libs = "/tmp/python3.11/kernel-libs/lib/python3.11/site-packages" + + venv.import_package_bundle( + server_libs, condition_env="DEEPNOTE_INCLUDE_SERVER_PACKAGES" + ) + venv.import_package_bundle(system_site, priority=True) + venv.import_package_bundle(kernel_libs) + + pth_file = Path(venv.site_packages_path) / "deepnote.pth" + lines = pth_file.read_text().splitlines() + + assert len(lines) == 3 + # Line 1: server-libs conditional + assert "DEEPNOTE_INCLUDE_SERVER_PACKAGES" in lines[0] + assert f"sys.path.insert(0, '{server_libs}')" in lines[0] + # Line 2: system with priority + assert f"sys.path.insert(0, '{system_site}')" in lines[1] + # Line 3: kernel plain path + assert lines[2].strip() == kernel_libs + + def test_import_package_bundle_condition_env_and_priority_mutually_exclusive( + self, venv + ): + """Test that condition_env and priority cannot be used together.""" + bundle_path = "/some/bundle/site-packages" + + with pytest.raises(ValueError, match="mutually exclusive"): + venv.import_package_bundle( + bundle_path, + condition_env="SOME_ENV_VAR", + priority=True, + )