|
| 1 | +"""Cloud test for the MicroVM session example that suspends after a job. |
| 2 | +
|
| 3 | +The local runner cannot launch a MicroVM. So the test is cloud-only. |
| 4 | +""" |
| 5 | + |
| 6 | +import uuid |
| 7 | + |
| 8 | +import pytest |
| 9 | +from aws_durable_execution_sdk_python.execution import InvocationStatus |
| 10 | +from aws_durable_execution_sdk_python.lambda_service import ( |
| 11 | + OperationStatus, |
| 12 | + OperationType, |
| 13 | +) |
| 14 | + |
| 15 | + |
| 16 | +# The examples package does not depend on the extras. So an environment |
| 17 | +# without the extras skips this module instead of failing to import the |
| 18 | +# example. |
| 19 | +pytest.importorskip("aws_durable_execution_sdk_python_extras") |
| 20 | + |
| 21 | +from src.microvm import microvm_session_suspend_after # noqa: E402 |
| 22 | +from test.conftest import deserialize_operation_payload # noqa: E402 |
| 23 | +from test.microvm.microvm_helpers import ( # noqa: E402 |
| 24 | + expect_microvm_terminated, |
| 25 | + fail_callback_after_test_failure, |
| 26 | + find_launched_microvm_id, |
| 27 | + wait_for_microvm_state, |
| 28 | +) |
| 29 | + |
| 30 | + |
| 31 | +def tree(operations): |
| 32 | + """Return (name, type, status, children) for each operation, in order.""" |
| 33 | + return [ |
| 34 | + ( |
| 35 | + operation.name, |
| 36 | + operation.operation_type, |
| 37 | + operation.status, |
| 38 | + tree(getattr(operation, "child_operations", None) or []), |
| 39 | + ) |
| 40 | + for operation in operations |
| 41 | + ] |
| 42 | + |
| 43 | + |
| 44 | +SUCCEEDED = OperationStatus.SUCCEEDED |
| 45 | + |
| 46 | + |
| 47 | +@pytest.mark.example |
| 48 | +@pytest.mark.durable_execution( |
| 49 | + handler=microvm_session_suspend_after.handler, |
| 50 | + lambda_function_name="MicroVM Session Suspend After", |
| 51 | + cloud_only=True, |
| 52 | +) |
| 53 | +def test_suspends_the_microvm_after_the_first_job_and_resumes_it_for_the_next( |
| 54 | + durable_runner, |
| 55 | +): |
| 56 | + key = str(uuid.uuid4()) |
| 57 | + |
| 58 | + with durable_runner: |
| 59 | + # A synchronous Invoke returns only when the execution ends. The cloud |
| 60 | + # runner's boto3 client reads for at most 60 seconds, and it sends the |
| 61 | + # Invoke again after a read timeout. Each resend starts a new |
| 62 | + # execution, which launches its own MicroVM. So the test starts the |
| 63 | + # execution asynchronously and polls for its result. |
| 64 | + execution_arn = durable_runner.run_async(input={"key": key}, timeout=300) |
| 65 | + |
| 66 | + # The test reads the MicroVM state while the gate is pending. So it |
| 67 | + # reads the MicroVM ID from the launch step in the execution history. |
| 68 | + microvm_id = find_launched_microvm_id( |
| 69 | + execution_arn, "suspend-after-session.launch" |
| 70 | + ) |
| 71 | + |
| 72 | + callback_id = None |
| 73 | + gate_completed = False |
| 74 | + try: |
| 75 | + # The cloud runner reads the callback ID from the CallbackStarted |
| 76 | + # event named gate. It returns the ID only after the invocation |
| 77 | + # that created the callback has completed. |
| 78 | + callback_id = durable_runner.wait_for_callback( |
| 79 | + execution_arn=execution_arn, name="gate", timeout=120 |
| 80 | + ) |
| 81 | + |
| 82 | + # 1. The gate starts after the first job ends. |
| 83 | + # 2. The second job starts only after the test completes the gate. |
| 84 | + # 3. The session does not suspend when idle. So only the first |
| 85 | + # job's suspend_after can suspend the MicroVM. |
| 86 | + # 4. So a SUSPENDED state read now lies between the two jobs, and |
| 87 | + # comes from suspend_after. |
| 88 | + # The poll lasts at most 150 seconds. The gate times out after 5 |
| 89 | + # minutes. So the gate stays pending for the whole poll. The |
| 90 | + # state can also pass through SUSPENDING. So the test requires |
| 91 | + # only that SUSPENDED is seen once. |
| 92 | + wait_for_microvm_state(microvm_id, "SUSPENDED") |
| 93 | + |
| 94 | + durable_runner.send_callback_success(callback_id=callback_id) |
| 95 | + gate_completed = True |
| 96 | + finally: |
| 97 | + # A failure before the gate completes leaves the execution waiting |
| 98 | + # for the gate, with the MicroVM still there. A callback failure |
| 99 | + # makes the session terminate the MicroVM now, instead of after |
| 100 | + # the gate's timeout. If the test fails before it has the callback |
| 101 | + # ID, the gate times out after 5 minutes. The session then |
| 102 | + # terminates the MicroVM. |
| 103 | + if callback_id is not None and not gate_completed: |
| 104 | + fail_callback_after_test_failure(durable_runner, callback_id) |
| 105 | + |
| 106 | + result = durable_runner.wait_for_result( |
| 107 | + execution_arn=execution_arn, timeout=300 |
| 108 | + ) |
| 109 | + |
| 110 | + assert result.status is InvocationStatus.SUCCEEDED, result.error |
| 111 | + output = deserialize_operation_payload(result.result) |
| 112 | + # 1. GetMicrovm reported the MicroVM SUSPENDED while the gate was pending. |
| 113 | + # 2. The second job read the file that the first job wrote. |
| 114 | + # 3. A suspended MicroVM keeps its files. |
| 115 | + # So both jobs ran in the same MicroVM, before and after the suspend. |
| 116 | + assert output["sessionMicrovmId"] == microvm_id |
| 117 | + assert output["writtenBy"] == microvm_id |
| 118 | + assert output["readBy"] == microvm_id |
| 119 | + assert output["value"] == f"written-by-{key}" |
| 120 | + |
| 121 | + expect_microvm_terminated(microvm_id) |
| 122 | + |
| 123 | + # The tree is the tree of the suspend example. The state check and the |
| 124 | + # resume run inside the second request step. So the tree shows them only |
| 125 | + # as a request step that succeeded after the suspend. |
| 126 | + assert tree(result.operations) == [ |
| 127 | + ( |
| 128 | + "suspend-after-session", |
| 129 | + OperationType.CONTEXT, |
| 130 | + SUCCEEDED, |
| 131 | + [ |
| 132 | + ("suspend-after-session.session", OperationType.STEP, SUCCEEDED, []), |
| 133 | + ("suspend-after-session.launch", OperationType.STEP, SUCCEEDED, []), |
| 134 | + ( |
| 135 | + "write", |
| 136 | + OperationType.CONTEXT, |
| 137 | + SUCCEEDED, |
| 138 | + [ |
| 139 | + ("write.callback", OperationType.CALLBACK, SUCCEEDED, []), |
| 140 | + ("write.request", OperationType.STEP, SUCCEEDED, []), |
| 141 | + ], |
| 142 | + ), |
| 143 | + ("gate", OperationType.CALLBACK, SUCCEEDED, []), |
| 144 | + ( |
| 145 | + "read", |
| 146 | + OperationType.CONTEXT, |
| 147 | + SUCCEEDED, |
| 148 | + [ |
| 149 | + ("read.callback", OperationType.CALLBACK, SUCCEEDED, []), |
| 150 | + ("read.request", OperationType.STEP, SUCCEEDED, []), |
| 151 | + ], |
| 152 | + ), |
| 153 | + ( |
| 154 | + "suspend-after-session.terminate", |
| 155 | + OperationType.STEP, |
| 156 | + SUCCEEDED, |
| 157 | + [], |
| 158 | + ), |
| 159 | + ], |
| 160 | + ) |
| 161 | + ] |
| 162 | + |
| 163 | + # The session sets neither auto_suspend_on_idle nor an idle policy. So the |
| 164 | + # launch records no resume mode. suspend_after is an option of one job, |
| 165 | + # and the launch checkpoint does not record it. |
| 166 | + session = result.get_context("suspend-after-session") |
| 167 | + launch = deserialize_operation_payload( |
| 168 | + session.get_step("suspend-after-session.launch").result |
| 169 | + ) |
| 170 | + assert launch["delivery"] == "http" |
| 171 | + assert launch["microvmId"] == microvm_id |
| 172 | + assert launch["resumeBeforeDelivery"] is False |
0 commit comments