Skip to content
2 changes: 1 addition & 1 deletion installer/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What kind of paths are in front here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The issue was with server site packages. Since they're dynamically inserted into sys path, they would always be prioritized.

venv.import_package_bundle(kernel_site_package_path)

# Phase 7.2: Symlink notebook to jupyter_server for compatibility with
Expand Down
17 changes: 15 additions & 2 deletions installer/module/virtual_environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,14 +81,18 @@ 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'.
:param priority: If True, insert at front of sys.path to override other bundles.
"""
pth_file_path = os.path.join(self._site_packages_path, "deepnote.pth")

Expand All @@ -105,6 +109,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(
Expand Down
91 changes: 91 additions & 0 deletions tests/unit/test_virtual_environment.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
"""Tests for VirtualEnvironment class."""

import os
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
os.makedirs(venv.site_packages_path, 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 = os.path.join(venv.site_packages_path, "deepnote.pth")
with open(pth_file, "r") as f:
content = f.read()

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 = os.path.join(venv.site_packages_path, "deepnote.pth")
with open(pth_file, "r") as f:
content = f.read()

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 = os.path.join(venv.site_packages_path, "deepnote.pth")
with open(pth_file, "r") as f:
content = f.read()

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 = os.path.join(venv.site_packages_path, "deepnote.pth")
with open(pth_file, "r") as f:
lines = f.readlines()

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