From dfe63613395ff53395b9b7fdcba9f14140c24553 Mon Sep 17 00:00:00 2001
From: technillogue <technillogue@gmail.com>
Date: Mon, 6 Nov 2023 16:48:47 -0500
Subject: [PATCH] review changes to tests

Signed-off-by: technillogue <technillogue@gmail.com>
---
 python/tests/server/test_http.py   | 7 +++----
 python/tests/server/test_runner.py | 5 +++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/tests/server/test_http.py b/python/tests/server/test_http.py
index 6f0862f072..5e2f8410d6 100644
--- a/python/tests/server/test_http.py
+++ b/python/tests/server/test_http.py
@@ -450,7 +450,6 @@ def test_prediction_idempotent_endpoint_conflict(client, match):
         json={"input": {"sleep": 1}},
         headers={"Prefer": "respond-async"},
     )
-    time.sleep(0.001)
     resp2 = client.put(
         "/predictions/5678efgh",
         json={"input": {"sleep": 1}},
@@ -518,12 +517,12 @@ def test_prediction_cancel(client):
     )
     assert resp.status_code == 202
 
-    resp = client.post("/predictions/123/cancel")
-    assert resp.status_code == 200
-
     resp = client.post("/predictions/456/cancel")
     assert resp.status_code == 404
 
+    resp = client.post("/predictions/123/cancel")
+    assert resp.status_code == 200
+
 
 @uses_predictor_with_client_options(
     "setup_weights",
diff --git a/python/tests/server/test_runner.py b/python/tests/server/test_runner.py
index ed59eab32d..30292f6383 100644
--- a/python/tests/server/test_runner.py
+++ b/python/tests/server/test_runner.py
@@ -75,7 +75,8 @@ async def test_prediction_runner_called_while_busy(runner):
 
     assert runner.is_busy()
     with pytest.raises(RunnerBusyError):
-        await runner.predict(request)[1]
+        _, task = runner.predict(request)
+        await task
 
     # Await to ensure that the first prediction is scheduled before we
     # attempt to shut down the runner.
@@ -90,7 +91,7 @@ async def test_prediction_runner_called_while_busy_idempotent(runner):
     runner.predict(request)
     _, async_result = runner.predict(request)
 
-    response = await async_result
+    response = await asyncio.wait_for(async_result, timeout=1)
     assert response.id == "abcd1234"
     assert response.output == "done in 0.1 seconds"
     assert response.status == "succeeded"