Added tests related to issue #1737

tests/unit/workflow/test_workflow.py

@@ -18,6 +18,7 @@
 import math
 import os
 import shutil
+import sys
 
 try:
     import cudf
@@ -667,3 +668,70 @@ def test_workflow_saved_schema(tmpdir):
     for node in postorder_iter_nodes(workflow2.output_node):
         assert node.input_schema is not None
         assert node.output_schema is not None
+
+def test_workflow_infer_modules_byvalue(tmp_path):
+    module_fn = tmp_path / "not_a_real_module.py"
+    sys.path.append(str(tmp_path))
+
+    with open(module_fn, "w") as module_f:
+        module_f.write("def identity(col):\n    return col")
+
+    import not_a_real_module
+
+    f_0 = not_a_real_module.identity
+    f_1 = lambda x: not_a_real_module.identity(x)
+    f_2 = lambda x: f_0(x)
+
+    try:
+        for fn, f in {"not_a_real_module.identity" : f_0, 
+                      "lambda x: not_a_real_module.identity(x)" : f_1, 
+                      "lambda x: f_0(x)" : f_2}.items():
+            assert(not_a_real_module in Workflow._getmodules([f]), f"inferred module dependencies from {fn}")
+
+    finally:
+        sys.path.pop()
+        del sys.modules["not_a_real_module"]
+
+
+def test_workflow_explicit_modules_byvalue(tmp_path):
+    module_fn = tmp_path / "not_a_real_module.py"
+    sys.path.append(str(tmp_path))
+
+    with open(module_fn, "w") as module_f:
+        module_f.write("def identity(col):\n    return col")
+
+    import not_a_real_module
+
+    wf = nvt.Workflow(
+        ["col_a"] >> nvt.ops.LambdaOp(not_a_real_module.identity)
+    )
+
+    wf.save(str(tmp_path / "identity-workflow"), modules_byvalue=[not_a_real_module])
+
+    del not_a_real_module
+    del sys.modules["not_a_real_module"]
+    os.unlink(str(tmp_path / "not_a_real_module.py"))
+
+    Workflow.load(str(tmp_path / "identity-workflow"))
+
+def test_workflow_auto_infer_modules_byvalue(tmp_path):
+    module_fn = tmp_path / "not_a_real_module.py"
+    sys.path.append(str(tmp_path))
+
+    with open(module_fn, "w") as module_f:
+        module_f.write("def identity(col):\n    return col")
+
+    import not_a_real_module
+
+    wf = nvt.Workflow(
+        ["col_a"] >> nvt.ops.LambdaOp(not_a_real_module.identity)
+    )
+
+    wf.save(str(tmp_path / "identity-workflow"), modules_byvalue="auto")
+
+    del not_a_real_module
+    del sys.modules["not_a_real_module"]
+    os.unlink(str(tmp_path / "not_a_real_module.py"))
+
+    Workflow.load(str(tmp_path / "identity-workflow"))
+