test: Fix cancellation test (#7296)

Signed-off-by: Graham King <grahamk@nvidia.com>

lib/bindings/python/tests/cancellation/test_cancellation.py

@@ -7,7 +7,11 @@ import pytest
 
 from dynamo.runtime import Context
 
-pytestmark = pytest.mark.pre_merge
+pytestmark = [
+    pytest.mark.gpu_0,
+    pytest.mark.pre_merge,
+    pytest.mark.unit,
+]
 
 
 class MockServer:
@@ -247,7 +251,9 @@ async def test_server_context_cancel(temp_file_store, server, client):
     except ValueError as e:
         # Verify the expected cancellation exception is received
         # TODO: Should this be a asyncio.CancelledError?
-        assert str(e).startswith("Stream ended before generation completed")
+        assert str(e).startswith(
+            "Disconnected: Stream ended before generation completed"
+        )
 
     # Verify server context cancellation status
     assert handler.context_is_stopped
@@ -272,9 +278,8 @@ async def test_server_raise_cancelled(temp_file_store, server, client):
     except ValueError as e:
         # Verify the expected cancellation exception is received
         # TODO: Should this be a asyncio.CancelledError?
-        assert (
-            str(e)
-            == "a python exception was caught while processing the async generator: CancelledError: "
+        assert str(e).endswith(
+            "a python exception was caught while processing the async generator: CancelledError: "
         )
 
     # Verify server context cancellation status