@@ -671,6 +671,57 @@ def test_heartbeat_starts_and_stops(self) -> None:
671671
672672 mock .extend_reservation .assert_called ()
673673
674+ def test_heartbeat_skipped_when_ttl_zero (self ) -> None :
675+ mock = _make_mock_client ()
676+ sr = StreamReservation (
677+ mock ,
678+ subject = _default_subject (),
679+ action = _default_action (),
680+ estimate = _default_estimate (),
681+ ttl_ms = 0 ,
682+ )
683+ assert sr ._start_heartbeat () is None
684+
685+ def test_heartbeat_extend_failure_logged (self ) -> None :
686+ mock = _make_mock_client ()
687+ mock .create_reservation .return_value = _allow_response ()
688+ mock .commit_reservation .return_value = _commit_success ()
689+ mock .extend_reservation .return_value = CyclesResponse .http_error (
690+ 500 , "Server error" , body = {"error" : "INTERNAL" },
691+ )
692+
693+ sr = StreamReservation (
694+ mock ,
695+ subject = _default_subject (),
696+ action = _default_action (),
697+ estimate = _default_estimate (),
698+ ttl_ms = 2000 ,
699+ )
700+
701+ with sr :
702+ time .sleep (1.2 )
703+
704+ mock .extend_reservation .assert_called ()
705+
706+ def test_heartbeat_extend_exception_logged (self ) -> None :
707+ mock = _make_mock_client ()
708+ mock .create_reservation .return_value = _allow_response ()
709+ mock .commit_reservation .return_value = _commit_success ()
710+ mock .extend_reservation .side_effect = RuntimeError ("boom" )
711+
712+ sr = StreamReservation (
713+ mock ,
714+ subject = _default_subject (),
715+ action = _default_action (),
716+ estimate = _default_estimate (),
717+ ttl_ms = 2000 ,
718+ )
719+
720+ with sr :
721+ time .sleep (1.2 )
722+
723+ mock .extend_reservation .assert_called ()
724+
674725
675726# ---------------------------------------------------------------------------
676727# AsyncStreamReservation tests
@@ -1281,6 +1332,77 @@ async def test_caps_propagated(self) -> None:
12811332 assert reservation .caps is not None
12821333 assert reservation .caps .max_tokens == 256
12831334
1335+ @pytest .mark .asyncio
1336+ async def test_decision_property (self ) -> None :
1337+ mock = _make_async_mock_client ()
1338+ mock .create_reservation .return_value = _allow_response ()
1339+ mock .commit_reservation .return_value = _commit_success ()
1340+
1341+ asr = AsyncStreamReservation (
1342+ mock ,
1343+ subject = _default_subject (),
1344+ action = _default_action (),
1345+ estimate = _default_estimate (),
1346+ ttl_ms = 1000 ,
1347+ )
1348+
1349+ async with asr as reservation :
1350+ assert reservation .decision is not None
1351+
1352+ @pytest .mark .asyncio
1353+ async def test_heartbeat_skipped_when_ttl_zero (self ) -> None :
1354+ mock = _make_async_mock_client ()
1355+ asr = AsyncStreamReservation (
1356+ mock ,
1357+ subject = _default_subject (),
1358+ action = _default_action (),
1359+ estimate = _default_estimate (),
1360+ ttl_ms = 0 ,
1361+ )
1362+ assert asr ._start_heartbeat () is None
1363+
1364+ @pytest .mark .asyncio
1365+ async def test_heartbeat_extend_failure_logged (self ) -> None :
1366+ mock = _make_async_mock_client ()
1367+ mock .create_reservation .return_value = _allow_response ()
1368+ mock .commit_reservation .return_value = _commit_success ()
1369+ mock .extend_reservation .return_value = CyclesResponse .http_error (
1370+ 500 , "Server error" , body = {"error" : "INTERNAL" },
1371+ )
1372+
1373+ asr = AsyncStreamReservation (
1374+ mock ,
1375+ subject = _default_subject (),
1376+ action = _default_action (),
1377+ estimate = _default_estimate (),
1378+ ttl_ms = 2000 ,
1379+ )
1380+
1381+ async with asr :
1382+ await asyncio .sleep (1.2 )
1383+
1384+ mock .extend_reservation .assert_called ()
1385+
1386+ @pytest .mark .asyncio
1387+ async def test_heartbeat_extend_exception_logged (self ) -> None :
1388+ mock = _make_async_mock_client ()
1389+ mock .create_reservation .return_value = _allow_response ()
1390+ mock .commit_reservation .return_value = _commit_success ()
1391+ mock .extend_reservation .side_effect = RuntimeError ("boom" )
1392+
1393+ asr = AsyncStreamReservation (
1394+ mock ,
1395+ subject = _default_subject (),
1396+ action = _default_action (),
1397+ estimate = _default_estimate (),
1398+ ttl_ms = 2000 ,
1399+ )
1400+
1401+ async with asr :
1402+ await asyncio .sleep (1.2 )
1403+
1404+ mock .extend_reservation .assert_called ()
1405+
12841406
12851407class TestBudgetExceeded :
12861408 def test_budget_exceeded_raises_typed_error (self ) -> None :
0 commit comments