From 286983f3eec1b1e4a73fa1ce0683c5b149c74a58 Mon Sep 17 00:00:00 2001 From: troychiu Date: Thu, 7 Dec 2023 10:25:01 -0800 Subject: [PATCH] fix tests Signed-off-by: troychiu --- .../flytekit-flyin/tests/test_flyin_plugin.py | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/plugins/flytekit-flyin/tests/test_flyin_plugin.py b/plugins/flytekit-flyin/tests/test_flyin_plugin.py index 9b1e1a5297..2088047862 100644 --- 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):