Fix unit-test

Signed-off-by: Mecoli1219 <[email protected]>

flytekit/remote/init_remote.py

@@ -3,8 +3,8 @@
 
 from flytekit.configuration import Config
 from flytekit.remote.remote import FlyteRemote
-from flytekit.tools.translator import Options
 from flytekit.tools.interactive import ipython_check
+from flytekit.tools.translator import Options
 
 REMOTE_ENTRY: typing.Optional[FlyteRemote] = None
 # TODO: This should be merged into the FlyteRemote in the future

tests/flytekit/integration/remote/test_remote.py

@@ -8,6 +8,7 @@
 
 import joblib
 import pytest
+from mock import mock, patch
 
 from flytekit import LaunchPlan, kwtypes
 from flytekit.configuration import Config, ImageConfig, SerializationSettings
@@ -16,7 +17,7 @@
 from flytekit.core.workflow import reference_workflow
 from flytekit.exceptions.user import FlyteAssertion, FlyteEntityNotExistException
 from flytekit.extras.sqlite3.task import SQLite3Config, SQLite3Task
-from flytekit.remote.remote import FlyteRemote
+from flytekit.remote import init_remote, FlyteRemote
 from flytekit.types.schema import FlyteSchema
 
 MODULE_PATH = pathlib.Path(__file__).parent / "workflows/basic"
@@ -25,7 +26,7 @@
 PROJECT = "flytesnacks"
 DOMAIN = "development"
 VERSION = f"v{os.getpid()}"
-
+init_remote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN, interactive_mode_enabled=True)
 
 @pytest.fixture(scope="session")
 def register():
@@ -483,3 +484,101 @@ def test_execute_workflow_with_maptask(register):
         wait=True,
     )
     assert execution.outputs["o0"] == [4, 5, 6]
+
+
+@mock.patch("flytekit.tools.interactive.ipython_check")
+def test_jupyter_remote_workflow(mock_ipython_check):
+    mock_ipython_check.return_value = True
+    with pytest.raises(AssertionError):
+        init_remote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN, interactive_mode_enabled=True)
+    from .workflows.basic.child_workflow import parent_wf, double
+
+    # child_workflow.parent_wf asynchronously register a parent wf1 with child lp from another wf2.
+    future0 = double.remote(a=3)
+    future1 = parent_wf.remote(a=3)
+    future2 = parent_wf.remote(a=2)
+    assert future0.version != VERSION
+    assert future1.version != VERSION
+    assert future2.version != VERSION
+    # It should generate a new version for each execution
+    assert future1.version != future2.version
+
+    out0 = future0.wait()
+    assert out0.outputs["o0"] == 6
+    out1 = future1.wait()
+    assert out1.outputs["o0"] == 18
+    out2 = future2.wait()
+    assert out2.outputs["o0"] == 12
+
+
+def test_execute_jupyter_python_task(register):
+    """Test execution of a @task-decorated python function that is already registered."""
+    with patch("flytekit.tools.interactive.ipython_check") as mock_ipython_check:
+        mock_ipython_check.return_value = True
+
+        from .workflows.basic.basic_workflow import t1
+
+        future = t1.remote(a=10, version=VERSION)
+        out = future.wait()
+        assert future.version == VERSION
+
+        assert out.outputs["t1_int_output"] == 12
+        assert out.outputs["c"] == "world"
+
+
+def test_execute_jupyter_python_workflow(register):
+    """Test execution of a @workflow-decorated python function."""
+    with patch("flytekit.tools.interactive.ipython_check") as mock_ipython_check:
+        mock_ipython_check.return_value = True
+        from .workflows.basic.basic_workflow import my_basic_wf
+
+        future = my_basic_wf.remote(a=10, b="xyz", version=VERSION)
+        out = future.wait()
+        assert out.outputs["o0"] == 12
+        assert out.outputs["o1"] == "xyzworld"
+
+
+@mock.patch("flytekit.tools.interactive.ipython_check")
+def test_fetch_execute_task_list_of_floats(mock_ipython_check):
+    mock_ipython_check.return_value = True
+    from .workflows.basic.list_float_wf import concat_list
+
+    xs: typing.List[float] = [0.1, 0.2, 0.3, 0.4, -99999.7]
+    future = concat_list.remote(xs=xs)
+    out = future.wait()
+    assert out.outputs["o0"] == "[0.1, 0.2, 0.3, 0.4, -99999.7]"
+
+
+@mock.patch("flytekit.tools.interactive.ipython_check")
+def test_fetch_execute_task_convert_dict(mock_ipython_check):
+    mock_ipython_check.return_value = True
+    from .workflows.basic.dict_str_wf import convert_to_string
+
+    d: typing.Dict[str, str] = {"key1": "value1", "key2": "value2"}
+    future = convert_to_string.remote(d=d)
+    out = future.wait()
+    assert json.loads(out.outputs["o0"]) == {"key1": "value1", "key2": "value2"}
+
+
+@mock.patch("flytekit.tools.interactive.ipython_check")
+def test_execute_jupyter_python_workflow_dict_of_string_to_string(mock_ipython_check):
+    """Test execution of a @workflow-decorated python function."""
+    mock_ipython_check.return_value = True
+    from .workflows.basic.dict_str_wf import my_dict_str_wf
+
+    d: typing.Dict[str, str] = {"k1": "v1", "k2": "v2"}
+    future = my_dict_str_wf.remote(d=d)
+    out = future.wait()
+    assert json.loads(out.outputs["o0"]) == {"k1": "v1", "k2": "v2"}
+
+
+@mock.patch("flytekit.tools.interactive.ipython_check")
+def test_execute_jupyter_python_workflow_list_of_floats(mock_ipython_check):
+    """Test execution of a @workflow-decorated python function."""
+    mock_ipython_check.return_value = True
+    from .workflows.basic.list_float_wf import my_list_float_wf
+
+    xs: typing.List[float] = [42.24, 999.1, 0.0001]
+    future = my_list_float_wf.remote(xs=xs)
+    out = future.wait()
+    assert out.outputs["o0"] == "[42.24, 999.1, 0.0001]"

tests/flytekit/integration/remote/workflows/basic/basic_workflow.py

@@ -22,7 +22,7 @@
 
 from flytekit import task, workflow
 
-IMAGE = os.getenv("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
+IMAGE = os.environ.get("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
 
 
 @task(container_image=IMAGE)

tests/flytekit/integration/remote/workflows/basic/child_workflow.py

@@ -3,7 +3,8 @@
 from flytekit import LaunchPlan, task, workflow
 from flytekit.models.common import Labels
 
-IMAGE = os.getenv("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
+IMAGE = os.environ.get("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
+
 
 @task(container_image=IMAGE)
 def double(a: int) -> int:
@@ -21,7 +22,9 @@ def my_childwf(a: int = 42) -> int:
     return b
 
 
-child_lp = LaunchPlan.get_or_create(my_childwf, name="my_fixed_child_lp", labels=Labels({"l1": "v1"}))
+child_lp = LaunchPlan.get_or_create(
+    my_childwf, name="my_fixed_child_lp", labels=Labels({"l1": "v1"})
+)
 
 
 @workflow

tests/flytekit/integration/remote/workflows/basic/dict_str_wf.py

@@ -4,7 +4,7 @@
 
 from flytekit import task, workflow
 
-IMAGE = os.getenv("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
+IMAGE = os.environ.get("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
 
 
 @task(container_image=IMAGE)
@@ -18,4 +18,6 @@ def my_dict_str_wf(d: typing.Dict[str, str]) -> str:
 
 
 if __name__ == "__main__":
-    print(f"Running my_wf(d={{'a': 'rwx', 'b': '42'}})={my_dict_str_wf(d={'a': 'rwx', 'b': '42'})}")
+    print(
+        f"Running my_wf(d={{'a': 'rwx', 'b': '42'}})={my_dict_str_wf(d={'a': 'rwx', 'b': '42'})}"
+    )

tests/flytekit/integration/remote/workflows/basic/list_float_wf.py

@@ -3,7 +3,7 @@
 
 from flytekit import task, workflow
 
-IMAGE = os.getenv("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
+IMAGE = os.environ.get("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
 
 
 @task(container_image=IMAGE)

tests/flytekit/unit/core/test_context_manager.py

@@ -267,7 +267,7 @@ def test_serialization_settings_transport():
     ss = SerializationSettings.from_transport(tp)
     assert ss is not None
     assert ss == serialization_settings
-    assert len(tp) == 408
+    assert len(tp) == 440
 
 
 def test_exec_params():