refactor logic

Author: Andrei Neagu

packages/service-library/tests/long_running_interfaces/test__rpc.py

@@ -136,6 +136,21 @@ async def server_rpc_interface(
     await server.teardown()
 
 
+async def _wait_to_be_finished(
+    client_rpc_interface: ClientRPCInterface, unique_id: JobUniqueId
+) -> None:
+    async for attempt in AsyncRetrying(
+        wait=wait_fixed(1),
+        stop=stop_after_delay(10),
+        retry=retry_if_exception_type(AssertionError),
+        reraise=True,
+    ):
+        with attempt:
+            assert (
+                await client_rpc_interface.get_status(unique_id) == JobStatus.FINISHED
+            )
+
+
 async def test_workflow(
     server_rpc_interface: ServerRPCInterface,
     client_rpc_interface: ClientRPCInterface,
@@ -157,23 +172,31 @@ async def test_workflow(
     with pytest.raises(NoResultIsAvailableError):
         assert await client_rpc_interface.get_result(unique_id)
 
-    # wait to be finsiehd
-    async for attempt in AsyncRetrying(
-        wait=wait_fixed(1),
-        stop=stop_after_delay(10),
-        retry=retry_if_exception_type(AssertionError),
-        reraise=True,
-    ):
-        with attempt:
-            assert (
-                await client_rpc_interface.get_status(unique_id) == JobStatus.FINISHED
-            )
+    await _wait_to_be_finished(client_rpc_interface, unique_id)
 
     # result should be ready
     assert await client_rpc_interface.get_result(unique_id) == ResultModel(
         data=f"{unique_id} done"
     )
 
+    # result and corresponding task is automatically removed once retruned
+    with pytest.raises(JobNotFoundError):
+        await client_rpc_interface.get_result(unique_id)
+
+
+async def test_finishes_with_error(
+    server_rpc_interface: ServerRPCInterface,
+    client_rpc_interface: ClientRPCInterface,
+    unique_id: JobUniqueId,
+    initilized_server_interface: _MockServerInterface,
+) -> None:
+
+    await client_rpc_interface.start(
+        "handler_name", unique_id, timeout=timedelta(seconds=4)
+    )
+
+    await _wait_to_be_finished(client_rpc_interface, unique_id)
+
     # simulates an error in the result
     initilized_server_interface.result_raises = True
     result = await client_rpc_interface.get_result(unique_id)

packages/service-library/tests/long_running_interfaces/test_workflow.py

@@ -211,21 +211,23 @@ async def _cancel_task_in_server(server: Server) -> None:
         task.cancel()
 
 
+async def _to_run(client: Client, retry_count: NonNegativeInt) -> None:
+    result = await client.ensure_result(
+        "sleep_for_f",
+        expected_type=type(None),
+        timeout=timedelta(seconds=10),
+        duration=2,
+        retry_count=retry_count,
+    )
+    assert result is None
+
+
 async def test_cancellation_from_server_retires_and_finishes(
     server: Server, client: Client
 ):
-    async def _to_run() -> None:
-        result = await client.ensure_result(
-            "sleep_for_f",
-            expected_type=type(None),
-            timeout=timedelta(seconds=10),
-            duration=2,
-        )
-        assert result is None
-
     await _assert_tasks_count(server, count=0)
 
-    task = asyncio.create_task(_to_run())
+    task = asyncio.create_task(_to_run(client, retry_count=3))
     await _assert_tasks_count(server, count=1)
 
     await _cancel_task_in_server(server)
@@ -236,19 +238,9 @@ async def _to_run() -> None:
 async def test_cancellation_from_server_fails_if_no_more_retries_available(
     server: Server, client: Client
 ):
-    async def _to_run() -> None:
-        result = await client.ensure_result(
-            "sleep_for_f",
-            expected_type=type(None),
-            timeout=timedelta(seconds=10),
-            duration=2,
-            retry_count=1,
-        )
-        assert result is None
-
     await _assert_tasks_count(server, count=0)
 
-    task = asyncio.create_task(_to_run())
+    task = asyncio.create_task(_to_run(client, retry_count=1))
     await _assert_tasks_count(server, count=1)
 
     await _cancel_task_in_server(server)