Test sending a Module directly to a new env.

tests/conftest.py

@@ -286,6 +286,7 @@ def event_loop():
     conda_package,  # noqa: F401
     git_package,  # noqa: F401
     installed_editable_package,  # noqa: F401
+    installed_editable_package_copy,  # noqa: F401
     local_package,  # noqa: F401
     package,  # noqa: F401
     pip_package,  # noqa: F401

tests/fixtures/package_fixtures.py

@@ -103,3 +103,28 @@ def installed_editable_package(tmp_path_factory):
 
     # Delete everything in tmp_package_dir recursively
     shutil.rmtree(tmp_package_dir)
+
+
+@pytest.fixture(scope="session")
+def installed_editable_package_copy(tmp_path_factory):
+    tmp_package_dir = (
+        tmp_path_factory.mktemp("fake_package_copy") / "test_fake_package_copy"
+    )
+
+    # Copy the test_fake_package directory that's in the same directory as this file, to the tmp_package_dir established
+    # above.
+    shutil.copytree(
+        os.path.join(os.path.dirname(__file__), "test_fake_package_copy"),
+        tmp_package_dir,
+    )
+
+    # Run a pip install -e on the tmp_package_dir via subprocess.run, locally, not on the cluster
+    subprocess.run(["pip", "install", "-e", str(tmp_package_dir)], check=True)
+
+    yield
+
+    # Uninstall the package after the test is done
+    subprocess.run(["pip", "uninstall", "-y", "test_fake_package_copy"], check=True)
+
+    # Delete everything in tmp_package_dir recursively
+    shutil.rmtree(tmp_package_dir)

tests/fixtures/test_fake_package/test_fake_package/__init__.py

@@ -1 +1,2 @@
 from .function_to_import import editable_package_function
+from .module_to_import import TestModuleFromPackage

tests/fixtures/test_fake_package/test_fake_package/module_to_import.py

@@ -0,0 +1,5 @@
+class TestModuleFromPackage:
+    @staticmethod
+    def hello_world():
+        print("Hello from the editable package module!")
+        return "Hello from the editable package module!"

tests/fixtures/test_fake_package_copy/setup.py

@@ -0,0 +1,12 @@
+from setuptools import find_packages, setup
+
+setup(
+    name="test_fake_package_copy",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=[],
+    author="Rohin Bhasin",
+    author_email="bhasin.rohin@gmail.com",
+    description="A simple example package",
+    python_requires=">=3.6",
+)

tests/fixtures/test_fake_package_copy/test_fake_package_copy/__init__.py

@@ -0,0 +1,2 @@
+from .function_to_import import editable_package_function
+from .module_to_import import TestModuleFromPackage

tests/fixtures/test_fake_package_copy/test_fake_package_copy/function_to_import.py

@@ -0,0 +1,3 @@
+def editable_package_function():
+    print("Hello from the editable package!")
+    return "Hello from the editable package!"

tests/fixtures/test_fake_package_copy/test_fake_package_copy/module_to_import.py

@@ -0,0 +1,5 @@
+class TestModuleFromPackage:
+    @staticmethod
+    def hello_world():
+        print("Hello from the editable package module!")
+        return "Hello from the editable package module!"

tests/test_resources/test_modules/test_functions/test_function.py

@@ -638,3 +638,9 @@ async def test_returns_coroutine(self, cluster):
         future_module = async_returns_coroutine_remote(2, 3, run_async=False)
         assert future_module.__class__.__name__ == "FutureModule"
         assert await future_module == 5
+
+    @pytest.mark.level("local")
+    def test_send_function_to_fresh_env(self, cluster):
+        env = rh.env(name="fresh_env", reqs=["numpy"])
+        summer_remote = rh.function(summer).to(cluster, env=env)
+        summer_remote(2, 3)

tests/test_resources/test_modules/test_module.py

@@ -918,3 +918,30 @@ def test_import_editable_package(self, cluster, installed_editable_package):
         # Now send this to the remote cluster and test that it can still be imported and used
         remote_editable_package_fn = rh.function(editable_package_function).to(cluster)
         assert remote_editable_package_fn() == "Hello from the editable package!"
+
+        cluster.delete("editable_package_function")
+        cluster.run("pip uninstall -y test_fake_package")
+
+    @pytest.mark.level("local")
+    def test_import_editable_package_from_new_env(
+        self, cluster, installed_editable_package_copy
+    ):
+        importlib_reload(site)
+
+        # Test that the editable package can be imported and used
+        from test_fake_package_copy import TestModuleFromPackage
+
+        assert (
+            TestModuleFromPackage.hello_world()
+            == "Hello from the editable package module!"
+        )
+
+        # Now send this to the remote cluster and test that it can still be imported and used
+        env = rh.env(name="fresh_env", reqs=["numpy"])
+        remote_editable_package_module = rh.module(TestModuleFromPackage).to(
+            cluster, env=env
+        )
+        assert (
+            remote_editable_package_module.hello_world()
+            == "Hello from the editable package module!"
+        )