fix tests

Signed-off-by: troychiu <[email protected]>
flyteorg · Dec 7, 2023 · 286983f · 286983f
1 parent a77058b
commit 286983f
Showing 1 changed file with 26 additions and 18 deletions.
diff --git a/plugins/flytekit-flyin/tests/test_flyin_plugin.py b/plugins/flytekit-flyin/tests/test_flyin_plugin.py
@@ -40,11 +40,6 @@ def mock_code_server_info_dict():
     return {"arm64": "Arm server info", "amd64": "AMD server info"}
 
 
-@pytest.fixture
-def mock_code_server_info_dict():
-    return {"arm64": "Arm server info", "amd64": "AMD server info"}
-
-
 @pytest.fixture
 def vscode_patches():
     with mock.patch("multiprocessing.Process") as mock_process, mock.patch(
@@ -54,15 +49,18 @@ def vscode_patches():
     ) as mock_exit_handler, mock.patch(
         "flytekitplugins.flyin.vscode_lib.decorator.download_vscode"
     ) as mock_download_vscode, mock.patch("signal.signal") as mock_signal, mock.patch(
-        "flytekitplugins.flyin.vscode_lib.decorator.generate_resume_task_script"
-    ) as mock_generate_resume_task_script:
+        "flytekitplugins.flyin.vscode_lib.decorator.prepare_resume_task_python"
+    ) as mock_prepare_resume_task_python, mock.patch(
+        "flytekitplugins.flyin.vscode_lib.decorator.prepare_launch_json"
+    ) as mock_prepare_launch_json:
         yield (
             mock_process,
             mock_prepare_interactive_python,
             mock_exit_handler,
             mock_download_vscode,
             mock_signal,
-            mock_generate_resume_task_script,
+            mock_prepare_resume_task_python,
+            mock_prepare_launch_json,
         )
 
 
@@ -73,7 +71,8 @@ def test_vscode_remote_execution(vscode_patches, mock_remote_execution):
         mock_exit_handler,
         mock_download_vscode,
         mock_signal,
-        mock_generate_resume_task_script,
+        mock_prepare_resume_task_python,
+        mock_prepare_launch_json,
     ) = vscode_patches
 
     @task
@@ -91,7 +90,8 @@ def wf():
     mock_exit_handler.assert_called_once()
     mock_prepare_interactive_python.assert_called_once()
     mock_signal.assert_called_once()
-    mock_generate_resume_task_script.assert_called_once()
+    mock_prepare_resume_task_python.assert_called_once()
+    mock_prepare_launch_json.assert_called_once()
 
 
 def test_vscode_remote_execution_but_disable(vscode_patches, mock_remote_execution):
@@ -101,7 +101,8 @@ def test_vscode_remote_execution_but_disable(vscode_patches, mock_remote_executi
         mock_exit_handler,
         mock_download_vscode,
         mock_signal,
-        mock_generate_resume_task_script,
+        mock_prepare_resume_task_python,
+        mock_prepare_launch_json,
     ) = vscode_patches
 
     @task
@@ -119,7 +120,8 @@ def wf():
     mock_exit_handler.assert_not_called()
     mock_prepare_interactive_python.assert_not_called()
     mock_signal.assert_not_called()
-    mock_generate_resume_task_script.assert_not_called()
+    mock_prepare_resume_task_python.assert_not_called()
+    mock_prepare_launch_json.assert_not_called()
 
 
 def test_vscode_local_execution(vscode_patches, mock_local_execution):
@@ -129,7 +131,8 @@ def test_vscode_local_execution(vscode_patches, mock_local_execution):
         mock_exit_handler,
         mock_download_vscode,
         mock_signal,
-        mock_generate_resume_task_script,
+        mock_prepare_resume_task_python,
+        mock_prepare_launch_json,
     ) = vscode_patches
 
     @task
@@ -147,7 +150,8 @@ def wf():
     mock_exit_handler.assert_not_called()
     mock_prepare_interactive_python.assert_not_called()
     mock_signal.assert_not_called()
-    mock_generate_resume_task_script.assert_not_called()
+    mock_prepare_resume_task_python.assert_not_called()
+    mock_prepare_launch_json.assert_not_called()
 
 
 def test_vscode_run_task_first_succeed(mock_remote_execution):
@@ -172,7 +176,8 @@ def test_vscode_run_task_first_fail(vscode_patches, mock_remote_execution):
         mock_exit_handler,
         mock_download_vscode,
         mock_signal,
-        mock_generate_resume_task_script,
+        mock_prepare_resume_task_python,
+        mock_prepare_launch_json,
     ) = vscode_patches
 
     @task
@@ -191,7 +196,8 @@ def wf(a: int, b: int):
     mock_exit_handler.assert_called_once()
     mock_prepare_interactive_python.assert_called_once()
     mock_signal.assert_called_once()
-    mock_generate_resume_task_script.assert_called_once()
+    mock_prepare_resume_task_python.assert_called_once()
+    mock_prepare_launch_json.assert_called_once()
 
 
 @mock.patch("flytekitplugins.flyin.jupyter_lib.decorator.subprocess.Popen")
@@ -254,7 +260,8 @@ def test_vscode_with_args(vscode_patches, mock_remote_execution):
         mock_exit_handler,
         mock_download_vscode,
         mock_signal,
-        mock_generate_resume_task_script,
+        mock_prepare_resume_task_python,
+        mock_prepare_launch_json,
     ) = vscode_patches
 
     @task
@@ -273,7 +280,8 @@ def wf():
     mock_exit_handler.assert_called_once()
     mock_prepare_interactive_python.assert_called_once()
     mock_signal.assert_called_once()
-    mock_generate_resume_task_script.assert_called_once()
+    mock_prepare_resume_task_python.assert_called_once()
+    mock_prepare_launch_json.assert_called_once()
 
 
 def test_vscode_extra_config(mock_remote_execution):