From 9b4e3118d4b12bc79eae48711d73faf1ac9b1894 Mon Sep 17 00:00:00 2001 From: rajdeep11 Date: Tue, 26 Nov 2024 12:45:05 +0530 Subject: [PATCH] optimising the get status Change-Id: I5fb886f8736a7b7f8fe2c553734957d6fcc22bd9 Signed-off-by: rajdeep11 --- tests/test_trainingmgr_util.py | 10 ++++++++-- trainingmgr/common/trainingmgr_util.py | 33 +++++++++++++++------------------ 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/tests/test_trainingmgr_util.py b/tests/test_trainingmgr_util.py index aa404bc..3a42c00 100644 --- a/tests/test_trainingmgr_util.py +++ b/tests/test_trainingmgr_util.py @@ -480,8 +480,14 @@ class Test_handle_async_feature_engineering_status_exception_case: class Test_get_one_word_status: def test_get_one_word_status(self): - steps_state = [0,1,2,3] - expected_data = "IN_PROGRESS" + steps_state = { + "DATA_EXTRACTION": "NOT_STARTED", + "DATA_EXTRACTION_AND_TRAINING": "NOT_STARTED", + "TRAINED_MODEL": "NOT_STARTED", + "TRAINING": "NOT_STARTED", + "TRAINING_AND_TRAINED_MODEL": "NOT_STARTED" + } + expected_data = "NOT_STARTED" assert get_one_word_status(steps_state) == expected_data,"data not equal" class Test_validate_trainingjob_name: diff --git a/trainingmgr/common/trainingmgr_util.py b/trainingmgr/common/trainingmgr_util.py index 083ccd7..61979d4 100644 --- a/trainingmgr/common/trainingmgr_util.py +++ b/trainingmgr/common/trainingmgr_util.py @@ -104,28 +104,25 @@ def check_key_in_dictionary(fields, dictionary): def get_one_word_status(steps_state): """ - This function converts steps_state to one word status(we call it overall_status also) - and return it. + Converts steps_state to a single word status (overall_status) and returns it. """ - failed_count = 0 - finished_count = 0 - not_started_count = 0 - in_progress_count = 0 - for step in steps_state: - if steps_state[step] == States.FAILED.name: - failed_count = failed_count + 1 - elif steps_state[step] == States.FINISHED.name: - finished_count = finished_count + 1 - elif steps_state[step] == States.NOT_STARTED.name: - not_started_count = not_started_count + 1 - else: - in_progress_count = in_progress_count + 1 - if failed_count > 0: + status_counts = { + States.FAILED.name: 0, + States.FINISHED.name: 0, + States.NOT_STARTED.name: 0, + States.IN_PROGRESS.name: 0, + } + + for step_status in steps_state.values(): + status_counts[step_status] += 1 + + if status_counts[States.FAILED.name] > 0: return States.FAILED.name - if not_started_count == len(steps_state): + if status_counts[States.NOT_STARTED.name] == len(steps_state): return States.NOT_STARTED.name - if finished_count == len(steps_state): + if status_counts[States.FINISHED.name] == len(steps_state): return States.FINISHED.name + return States.IN_PROGRESS.name -- 2.16.6