diff --git a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/__tests__/resquests.spec.ts b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/__tests__/resquests.spec.ts index 1cf6ec376e..73ef775906 100644 --- a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/__tests__/resquests.spec.ts +++ b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/__tests__/resquests.spec.ts @@ -13,6 +13,7 @@ import { jobRequest, jobRunRequest } from '../serverRequests'; +import { Dialog, showErrorMessage } from '@jupyterlab/apputils'; jest.mock('../handler', () => { return { @@ -20,6 +21,13 @@ jest.mock('../handler', () => { }; }); +jest.mock('@jupyterlab/apputils', () => ({ + showErrorMessage: jest.fn(), + Dialog: { + okButton: jest.fn() + } +})); + describe('jobdDataRequest', () => { afterEach(() => { jest.clearAllMocks(); @@ -213,6 +221,58 @@ describe('jobRequest()', () => { isSuccessful: false }); }); + + it('should show an error message if a task fails', async () => { + const mockData = { + [VdkOption.NAME]: 'Test Job', + [VdkOption.TEAM]: 'Test Team' + }; + + jobData.set(VdkOption.NAME, mockData[VdkOption.NAME]); + jobData.set(VdkOption.TEAM, mockData[VdkOption.TEAM]); + + const endpoint = 'DEPLOY'; + const taskId = endpoint + '-6266cd99-908c-480b-9a3e-8a30564736a4'; + const taskInitiationResponse = { + error: '', + message: `Task ${taskId} started` + }; + const taskCompletionResponse = { + task_id: taskId, + status: 'failed', + message: '', + error: 'An error occurred' + }; + + (requestAPI as jest.Mock) + .mockResolvedValueOnce(taskInitiationResponse) + .mockResolvedValue(taskCompletionResponse); + + const result = await jobRequest(endpoint); + + // Verify the call for initiating the task + expect(requestAPI).toHaveBeenCalledWith(endpoint, { + body: JSON.stringify(getJobDataJsonObject()), + method: 'POST' + }); + + // Verify the polling for task status + expect(requestAPI).toHaveBeenCalledWith(`taskStatus?taskId=${taskId}`, { + method: 'GET' + }); + + expect(showErrorMessage).toHaveBeenCalledWith( + 'Encountered an error while trying to connect the server. Error:', + taskCompletionResponse.error, + [Dialog.okButton()] + ); + + // Verify the final result + expect(result).toEqual({ + message: taskCompletionResponse.error, + isSuccessful: false + }); + }); }); describe('getFailingNotebookInfo', () => { diff --git a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/serverRequests.ts b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/serverRequests.ts index a9fb76cbc4..5e08e92bc5 100644 --- a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/serverRequests.ts +++ b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/src/serverRequests.ts @@ -104,11 +104,13 @@ const pollForTaskCompletion = async ( `taskStatus?taskId=${taskId}`, { method: 'GET' } ); - if ( - result.task_id === taskId && - (result.status === 'completed' || result.error) - ) { - return result; + if (result.task_id === taskId) { + if (result.status !== 'running') { + if (result.status === 'failed') { + showError(result.error); + } + return result; + } } } catch (error) { showError(error); @@ -151,10 +153,14 @@ export async function jobRunRequest(): Promise { const taskId = extractTaskIdFromMessage(initialResponse.message); const finalResult = await pollForTaskCompletion(taskId); - return { - message: finalResult.message as string, - isSuccessful: !finalResult.error - }; + if (finalResult.error) { + return { message: finalResult.error, isSuccessful: false }; + } else { + return { + message: finalResult.message as string, + isSuccessful: true + }; + } } catch (error) { showError(error); return { message: '', isSuccessful: false }; @@ -191,10 +197,14 @@ export async function jobRequest(endPoint: string): Promise { const taskId = extractTaskIdFromMessage(initialResponse.message); const finalResult = await pollForTaskCompletion(taskId); - return { - message: finalResult.message as string, - isSuccessful: !finalResult.error - }; + if (finalResult.error) { + return { message: finalResult.error, isSuccessful: false }; + } else { + return { + message: finalResult.message as string, + isSuccessful: true + }; + } } catch (error) { showError(error); return { message: '', isSuccessful: false }; diff --git a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/handlers.py b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/handlers.py index 481f1af1c4..86bcd1ce4c 100644 --- a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/handlers.py +++ b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/handlers.py @@ -51,20 +51,23 @@ def get(self): self.set_status(400) self.finish(json.dumps({"error": "taskId not provided."})) return - current_status = task_runner.get_status() - if current_status["task_id"] != task_id: - self.finish( - json.dumps( - { - "status": "failed", - "message": "Mismatched taskId.", - "error": f"Requested status for {task_id} but currently processing {current_status['task_id']}", - } + try: + current_status = task_runner.get_status() + if current_status["task_id"] != task_id: + self.finish( + json.dumps( + { + "status": "failed", + "message": "Mismatched taskId.", + "error": f"Requested status for {task_id} but currently processing {current_status['task_id']}", + } + ) ) - ) - return + return - self.finish(json.dumps(current_status)) + self.finish(json.dumps(current_status)) + except Exception as e: + self.finish(json.dumps({"message": f"{e}", "error": "true"})) class LoadJobDataHandler(APIHandler): @@ -114,18 +117,21 @@ class RunJobHandler(APIHandler): @tornado.web.authenticated def post(self): input_data = self.get_json_body() - task_id = task_runner.start_task( - "RUN", - lambda: VdkUI.run_job( - input_data[VdkOption.PATH.value], - input_data[VdkOption.ARGUMENTS.value], - ), - ) + try: + task_id = task_runner.start_task( + "RUN", + lambda: VdkUI.run_job( + input_data[VdkOption.PATH.value], + input_data[VdkOption.ARGUMENTS.value], + ), + ) - if task_id: - self.finish(task_start_response_success(task_id)) - else: - self.finish(task_start_response_failure("RUN")) + if task_id: + self.finish(task_start_response_success(task_id)) + else: + self.finish(task_start_response_failure("RUN")) + except Exception as e: + self.finish(json.dumps({"message": f"{e}", "error": "true"})) class DownloadJobHandler(APIHandler): @@ -140,19 +146,22 @@ class DownloadJobHandler(APIHandler): @tornado.web.authenticated def post(self): input_data = self.get_json_body() - task_id = task_runner.start_task( - "DOWNLOAD", - lambda: VdkUI.download_job( - input_data[VdkOption.NAME.value], - input_data[VdkOption.TEAM.value], - input_data[VdkOption.PATH.value], - ), - ) + try: + task_id = task_runner.start_task( + "DOWNLOAD", + lambda: VdkUI.download_job( + input_data[VdkOption.NAME.value], + input_data[VdkOption.TEAM.value], + input_data[VdkOption.PATH.value], + ), + ) - if task_id: - self.finish(task_start_response_success(task_id)) - else: - self.finish(task_start_response_failure("DOWNLOAD")) + if task_id: + self.finish(task_start_response_success(task_id)) + else: + self.finish(task_start_response_failure("DOWNLOAD")) + except Exception as e: + self.finish(json.dumps({"message": f"{e}", "error": "true"})) class ConvertJobHandler(APIHandler): @@ -164,15 +173,18 @@ class ConvertJobHandler(APIHandler): @tornado.web.authenticated def post(self): input_data = self.get_json_body() - task_id = task_runner.start_task( - "CONVERTJOBTONOTEBOOK", - lambda: VdkUI.convert_job(input_data[VdkOption.PATH.value]), - ) + try: + task_id = task_runner.start_task( + "CONVERTJOBTONOTEBOOK", + lambda: VdkUI.convert_job(input_data[VdkOption.PATH.value]), + ) - if task_id: - self.finish(task_start_response_success(task_id)) - else: - self.finish(task_start_response_failure("CONVERTJOBTONOTEBOOK")) + if task_id: + self.finish(task_start_response_success(task_id)) + else: + self.finish(task_start_response_failure("CONVERTJOBTONOTEBOOK")) + except Exception as e: + self.finish(json.dumps({"message": f"{e}", "error": "true"})) class CreateJobHandler(APIHandler): @@ -188,19 +200,22 @@ class CreateJobHandler(APIHandler): @tornado.web.authenticated def post(self): input_data = self.get_json_body() - task_id = task_runner.start_task( - "CREATE", - lambda: VdkUI.create_job( - input_data[VdkOption.NAME.value], - input_data[VdkOption.TEAM.value], - input_data[VdkOption.PATH.value], - ), - ) + try: + task_id = task_runner.start_task( + "CREATE", + lambda: VdkUI.create_job( + input_data[VdkOption.NAME.value], + input_data[VdkOption.TEAM.value], + input_data[VdkOption.PATH.value], + ), + ) - if task_id: - self.finish(task_start_response_success(task_id)) - else: - self.finish(task_start_response_failure("CREATE")) + if task_id: + self.finish(task_start_response_success(task_id)) + else: + self.finish(task_start_response_failure("CREATE")) + except Exception as e: + self.finish(json.dumps({"message": f"{e}", "error": "true"})) class CreateDeploymentHandler(APIHandler): @@ -215,20 +230,23 @@ class CreateDeploymentHandler(APIHandler): @tornado.web.authenticated def post(self): input_data = self.get_json_body() - task_id = task_runner.start_task( - "DEPLOY", - lambda: VdkUI.create_deployment( - input_data[VdkOption.NAME.value], - input_data[VdkOption.TEAM.value], - input_data[VdkOption.PATH.value], - input_data[VdkOption.DEPLOYMENT_REASON.value], - ), - ) + try: + task_id = task_runner.start_task( + "DEPLOY", + lambda: VdkUI.create_deployment( + input_data[VdkOption.NAME.value], + input_data[VdkOption.TEAM.value], + input_data[VdkOption.PATH.value], + input_data[VdkOption.DEPLOYMENT_REASON.value], + ), + ) - if task_id: - self.finish(task_start_response_success(task_id)) - else: - self.finish(task_start_response_failure("DEPLOY")) + if task_id: + self.finish(task_start_response_success(task_id)) + else: + self.finish(task_start_response_failure("DEPLOY")) + except Exception as e: + self.finish(json.dumps({"message": f"{e}", "error": "true"})) class GetNotebookInfoHandler(APIHandler): diff --git a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/task_runner.py b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/task_runner.py index d1876c089d..e1e5b94e09 100644 --- a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/task_runner.py +++ b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/task_runner.py @@ -41,7 +41,7 @@ def start_task(self, task_type, task_handler): """ task_id = f"{task_type}-{str(uuid.uuid4())}" with self.lock: - if self.__task_status["status"] not in ["idle", "completed"]: + if self.__task_status["status"] not in ["idle", "completed", "failed"]: return None self.__task_status = { @@ -51,9 +51,9 @@ def start_task(self, task_type, task_handler): "error": None, } - thread = threading.Thread(target=self._run_task, args=(task_handler,)) - thread.start() - return task_id + thread = threading.Thread(target=self._run_task, args=(task_handler,)) + thread.start() + return task_id def get_status(self): """ diff --git a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/tests/test_task_runner.py b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/tests/test_task_runner.py index 377c5f58c0..86868a0532 100644 --- a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/tests/test_task_runner.py +++ b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/tests/test_task_runner.py @@ -82,3 +82,25 @@ def test_task_status_after_completion(task_runner): assert ( task_runner.start_task(task_type, task_handler) is not None ), "Unable to start a new task after completion" + + +def test_task_status_fail_then_success(task_runner): + task_handler = Mock( + side_effect=[Exception("An error occurred"), "Task completed successfully"] + ) + task_type = "TEST" + task_id = task_runner.start_task(task_type, task_handler) + assert task_id is not None + + time.sleep(0.1) + status = task_runner.get_status() + assert status["status"] == "failed" + assert "An error occurred" in status["error"] + + task_id = task_runner.start_task(task_type, task_handler) + assert task_id is not None + + time.sleep(0.1) + status = task_runner.get_status() + assert status["status"] == "completed" + assert status["message"] == "Task completed successfully"