diff --git a/lib/redfish_code_update_utils.robot b/lib/redfish_code_update_utils.robot
index 916540e..4395ef2 100644
--- a/lib/redfish_code_update_utils.robot
+++ b/lib/redfish_code_update_utils.robot
@@ -341,6 +341,26 @@
     [Return]  ${task_inv}
 
 
+Get Task Inventory
+    [Documentation]  Return task inventory.
+    [Arguments]  ${task_info}
+
+    # Description of argument(s):
+    # task_info    Task information.
+
+    # Task information.
+    # @odata.id: /redfish/v1/TaskService/Tasks/1
+    # Id: 1
+    # TaskState: Starting
+    # TaskStatus: OK
+
+    ${task_payload}=  Redfish.Get Properties   ${task_info['@odata.id']}
+
+    ${task_inv}=  Create Initiated Task State Dict  ${task_payload}
+
+    [Return]  ${task_inv}
+
+
 Match Target URI
     [Documentation]  Match target uri from task list.
     [Arguments]  ${task_list}  ${target_uri}
diff --git a/lib/rest_client.robot b/lib/rest_client.robot
index c82087f..2ebbd14 100644
--- a/lib/rest_client.robot
+++ b/lib/rest_client.robot
@@ -372,7 +372,7 @@
     Valid Value  ret.status_code  ${valid_status_codes}
     Delete All Sessions
 
-    [Return]  ${ret.status_code}
+    [Return]  ${ret}
 
 
 Redfish Login
diff --git a/redfish/update_service/test_redfish_bmc_code_update.robot b/redfish/update_service/test_redfish_bmc_code_update.robot
index b02441e..5f6741d 100644
--- a/redfish/update_service/test_redfish_bmc_code_update.robot
+++ b/redfish/update_service/test_redfish_bmc_code_update.robot
@@ -82,7 +82,7 @@
     [Template]  Redfish Multiple Upload Image And Check Progress State
 
     # policy   image_file_path     alternate_image_file_path
-    Immediate  ${IMAGE_FILE_PATH}  ${ALTERNATE_IMAGE_FILE_PATH}
+    OnReset  ${IMAGE_FILE_PATH}  ${ALTERNATE_IMAGE_FILE_PATH}
 
 
 Post BMC Reset Perform Redfish Code Update
@@ -227,6 +227,7 @@
     ${redfish_update_uri}=  Get Redfish Update Service URI
 
     ${resp}=  Upload Image To BMC  ${redfish_update_uri}  timeout=${600}  data=${file_bin_data}
+
     Log To Console   Completed image upload to BMC.
 
     Sleep  5s
@@ -253,8 +254,13 @@
     # IMAGE_FILE_PATH            The path to BMC image file.
     # ALTERNATE_IMAGE_FILE_PATH  The path to alternate BMC image file.
 
+
+    ${task_inv_dict}=  Get Task State from File
+
     ${post_code_update_actions}=  Get Post Boot Action
+
     Valid File Path  ALTERNATE_IMAGE_FILE_PATH
+
     ${state}=  Get Pre Reboot State
     Rprint Vars  state
 
@@ -264,25 +270,35 @@
     # "HttpPushUri": "/redfish/v1/UpdateService/update",
 
     ${redfish_update_uri}=  Get Redfish Update Service URI
-    Redfish Upload Image  ${redfish_update_uri}  ${IMAGE_FILE_PATH}
 
-    ${first_image_id}=  Get Latest Image ID
-    Rprint Vars  first_image_id
-    Sleep  5s
-    Redfish Upload Image  ${redfish_update_uri}  ${ALTERNATE_IMAGE_FILE_PATH}
+    ${file_bin_data1}=  OperatingSystem.Get Binary File  ${IMAGE_FILE_PATH}
+    ${file_bin_data2}=  OperatingSystem.Get Binary File  ${ALTERNATE_IMAGE_FILE_PATH}
 
-    ${second_image_id}=  Get Latest Image ID
-    Rprint Vars  second_image_id
+    Log To Console  Uploading first image.
+    ${resp1}=  Upload Image To BMC  ${redfish_update_uri}  timeout=${600}  data=${file_bin_data1}
 
-    Check Image Update Progress State
-    ...  match_state='Updating', 'Disabled'  image_id=${second_image_id}
+    Log To Console  Uploading second image.
+    ${resp2}=  Upload Image To BMC  ${redfish_update_uri}  timeout=${600}  data=${file_bin_data2}
 
-    Check Image Update Progress State
-    ...  match_state='Updating'  image_id=${first_image_id}
+    ${task_info2}=    evaluate    json.loads('''${resp2.content}''')    json
 
-    Wait Until Keyword Succeeds  8 min  20 sec
-    ...  Check Image Update Progress State
-    ...    match_state='Enabled'  image_id=${first_image_id}
+    Sleep  3s
+
+    ${task_inv2}=  Get Task Inventory  ${task_info2}
+    Log  ${task_inv2}   
+
+    Wait Until Keyword Succeeds  5 min  10 sec
+    ...  Verify Task Progress State  ${task_inv2}  ${task_inv_dict['TaskException']}
+
+    ${task_info1}=    evaluate    json.loads('''${resp1.content}''')    json
+    Log  ${task_info1}
+
+    ${task_inv1}=  Get Task Inventory  ${task_info1}
+    Log  ${task_inv1}
+
+    Wait Until Keyword Succeeds  5 min  10 sec
+    ...  Verify Task Progress State  ${task_inv1}  ${task_inv_dict['TaskCompleted']}
+
     Run Key  ${post_code_update_actions['BMC image']['${apply_time}']}
     Redfish.Login
     Redfish Verify BMC Version  ${IMAGE_FILE_PATH}
diff --git a/redfish/update_service/test_redfish_image_upload.robot b/redfish/update_service/test_redfish_image_upload.robot
index ebb6936..19b4436 100644
--- a/redfish/update_service/test_redfish_image_upload.robot
+++ b/redfish/update_service/test_redfish_image_upload.robot
@@ -142,13 +142,13 @@
     Valid File Path  image_file_path
     Set ApplyTime  policy=OnReset
     ${image_data}=  OperatingSystem.Get Binary File  ${image_file_path}
-    ${status_code}=  Upload Image To BMC
+    ${resp}=  Upload Image To BMC
     ...  ${REDFISH_BASE_URI}UpdateService
     ...  ${timeout}
     ...  valid_status_codes=[${HTTP_OK}, ${HTTP_INTERNAL_SERVER_ERROR}]
     ...  data=${image_data}
 
-    Return From Keyword If  ${status_code} == ${HTTP_INTERNAL_SERVER_ERROR}
+    Return From Keyword If  ${resp.status_code} == ${HTTP_INTERNAL_SERVER_ERROR}
 
     ${image_id}=  Get Latest Image ID
     Rprint Vars  image_id