Fix prepare_payload_for_reinvocation in invoke (#1027)

Author: mircealam

src/rpdk/core/invoke.py

@@ -76,6 +76,8 @@ def get_contract_client(args, project):
 def prepare_payload_for_reinvocation(payload, response, artifact_type):
     if artifact_type == ARTIFACT_TYPE_RESOURCE:
         payload["callbackContext"] = response.get("callbackContext")
+        if ("resourceModel" in response) and ("requestData" in payload):
+            payload["requestData"]["resourceProperties"] = response.get("resourceModel")
 
     return payload
 

tests/test_invoke.py

@@ -8,7 +8,7 @@
 
 from rpdk.core.cli import EXIT_UNHANDLED_EXCEPTION, main
 from rpdk.core.contract.interface import Action, HookInvocationPoint
-from rpdk.core.invoke import _needs_reinvocation
+from rpdk.core.invoke import _needs_reinvocation, prepare_payload_for_reinvocation
 from rpdk.core.project import ARTIFACT_TYPE_HOOK, ARTIFACT_TYPE_RESOURCE, Project
 
 ACTIONS = list(Action.__members__)
@@ -351,6 +351,32 @@ def test_needs_reinvocation(max_reinvoke, current_invocation, result):
     assert _needs_reinvocation(max_reinvoke, current_invocation) is result
 
 
+# Test prepare_payload_for_reinvocation directly here
+def test_prepare_payload_for_reinvocation():
+    assert prepare_payload_for_reinvocation(
+        {}, {"callbackContext": {"foo": "bar"}}, ARTIFACT_TYPE_RESOURCE
+    ) == {"callbackContext": {"foo": "bar"}}
+    assert prepare_payload_for_reinvocation(
+        {"requestData": {}}, {"callbackContext": {"foo": "bar"}}, ARTIFACT_TYPE_RESOURCE
+    ) == {"callbackContext": {"foo": "bar"}, "requestData": {}}
+    assert prepare_payload_for_reinvocation(
+        {"requestData": {}},
+        {"callbackContext": {"foo": "bar"}, "resourceModel": {"foz": "baz"}},
+        ARTIFACT_TYPE_RESOURCE,
+    ) == {
+        "callbackContext": {"foo": "bar"},
+        "requestData": {"resourceProperties": {"foz": "baz"}},
+    }
+    assert prepare_payload_for_reinvocation(
+        {"requestData": {"resourceProperties": {}}},
+        {"callbackContext": {"foo": "bar"}, "resourceModel": {"foz": "baz"}},
+        ARTIFACT_TYPE_RESOURCE,
+    ) == {
+        "callbackContext": {"foo": "bar"},
+        "requestData": {"resourceProperties": {"foz": "baz"}},
+    }
+
+
 def _invoke_and_expect_resource(status, resource_payload_path, command, *args):
     (
         mock_project,