@@ -249,3 +249,353 @@ async def generator_clean_total() -> types.AsyncGenerator[int, None]:
249249 pass
250250
251251 assert i == 4
252+
253+
254+ @pytest .mark .asyncio
255+ async def test_aio_generator_timeout_detector_sync_callback (
256+ fake_clock : clock .FakeClock ,
257+ ) -> None :
258+ """Accept an ``on_timeout`` callback that is a plain function."""
259+ exceptions : types .List [BaseException ] = []
260+
261+ def on_timeout (
262+ generator : types .AsyncGenerator [int , None ],
263+ timeout : types .Optional [types .delta_type ],
264+ total_timeout : types .Optional [types .delta_type ],
265+ exception : BaseException ,
266+ ) -> None :
267+ """Record the timeout and return nothing to await."""
268+ exceptions .append (exception )
269+
270+ items : types .List [int ] = [
271+ i
272+ async for i in python_utils .aio_generator_timeout_detector (
273+ ticking_generator (fake_clock ),
274+ total_timeout = 0.45 ,
275+ on_timeout = on_timeout ,
276+ )
277+ ]
278+
279+ assert items == [0 , 1 , 2 , 3 , 4 ]
280+ assert len (exceptions ) == 1
281+ assert isinstance (exceptions [0 ], asyncio .TimeoutError )
282+
283+
284+ @pytest .mark .asyncio
285+ async def test_aio_generator_timeout_detector_async_callback (
286+ fake_clock : clock .FakeClock ,
287+ ) -> None :
288+ """Await an ``on_timeout`` callback that is a coroutine function."""
289+ calls : types .List [types .Mapping [str , types .Any ]] = []
290+
291+ async def on_timeout (
292+ generator : types .AsyncGenerator [int , None ],
293+ timeout : types .Optional [types .delta_type ],
294+ total_timeout : types .Optional [types .delta_type ],
295+ exception : BaseException ,
296+ ** kwargs : types .Mapping [str , types .Any ],
297+ ) -> None :
298+ """Record the extra keyword arguments after a real await."""
299+ await asyncio .sleep (0 )
300+ calls .append (kwargs )
301+
302+ items : types .List [int ] = [
303+ i
304+ async for i in python_utils .aio_generator_timeout_detector (
305+ ticking_generator (fake_clock ),
306+ total_timeout = 0.45 ,
307+ on_timeout = on_timeout ,
308+ context = {'attempt' : 1 },
309+ )
310+ ]
311+
312+ assert items == [0 , 1 , 2 , 3 , 4 ]
313+ assert calls == [{'context' : {'attempt' : 1 }}]
314+
315+
316+ async def collect_before_stall (
317+ generator : types .AsyncGenerator [int , None ],
318+ ) -> types .List [int ]:
319+ """Collect the items of a detector that has to stop at the stall."""
320+ items : types .List [int ] = []
321+
322+ async def collect () -> None :
323+ """Gather every item the detector lets through."""
324+ item : int
325+ async for item in generator :
326+ items .append (item )
327+
328+ try :
329+ # Without the guard a detector that waits out the stall would keep
330+ # the test busy for `STALL` seconds before it fails.
331+ await asyncio .wait_for (collect (), STALL / 10 )
332+ finally :
333+ await generator .aclose ()
334+
335+ return items
336+
337+
338+ @pytest .mark .asyncio
339+ async def test_aio_generator_timeout_detector_total_timeout_stall () -> None :
340+ """Enforce the total timeout while waiting for the next item."""
341+ items : types .List [int ] = await collect_before_stall (
342+ python_utils .aio_generator_timeout_detector (
343+ stalling_generator (), total_timeout = STALL / 500 , on_timeout = None
344+ )
345+ )
346+
347+ # The stall comes before item 5, so that item can never arrive in time.
348+ # How many of the earlier items arrive depends on the machine.
349+ assert items == list (range (len (items )))
350+ assert len (items ) <= 5
351+
352+
353+ @pytest .mark .asyncio
354+ async def test_aio_generator_timeout_detector_total_timeout_reraise () -> None :
355+ """Raise the total timeout error from a wait that it cut short."""
356+ detector : types .AsyncGenerator [int , None ] = (
357+ python_utils .aio_generator_timeout_detector (
358+ stalling_generator (), total_timeout = STALL / 500
359+ )
360+ )
361+
362+ with pytest .raises (asyncio .TimeoutError , match = 'Total timeout reached' ):
363+ await collect_before_stall (detector )
364+
365+
366+ @pytest .mark .asyncio
367+ async def test_aio_generator_timeout_detector_total_before_item () -> None :
368+ """Let the total timeout win when it ends before the item timeout."""
369+ detector : types .AsyncGenerator [int , None ] = (
370+ python_utils .aio_generator_timeout_detector (
371+ stalling_generator (), timeout = STALL , total_timeout = STALL / 500
372+ )
373+ )
374+
375+ with pytest .raises (asyncio .TimeoutError , match = 'Total timeout reached' ):
376+ await collect_before_stall (detector )
377+
378+
379+ @pytest .mark .asyncio
380+ async def test_aio_generator_timeout_detector_item_before_total () -> None :
381+ """Let the item timeout win when it ends before the total timeout."""
382+ detector : types .AsyncGenerator [int , None ] = (
383+ python_utils .aio_generator_timeout_detector (
384+ stalling_generator (), timeout = STALL / 500 , total_timeout = STALL
385+ )
386+ )
387+
388+ with pytest .raises (asyncio .TimeoutError ) as exc_info :
389+ await collect_before_stall (detector )
390+
391+ assert 'Total timeout reached' not in str (exc_info .value )
392+
393+
394+ @pytest .mark .asyncio
395+ async def test_aio_generator_timeout_detector_total_timeout_finishes () -> None :
396+ """Finish as normal when the generator ends within the total timeout."""
397+
398+ async def generator () -> types .AsyncGenerator [int , None ]:
399+ """Yield three items without waiting."""
400+ i : int
401+ for i in range (3 ):
402+ yield i
403+
404+ items : types .List [int ] = [
405+ i
406+ async for i in python_utils .aio_generator_timeout_detector (
407+ generator (), total_timeout = STALL
408+ )
409+ ]
410+
411+ assert items == [0 , 1 , 2 ]
412+
413+
414+ @pytest .mark .asyncio
415+ async def test_aio_generator_timeout_detector_generator_timeout () -> None :
416+ """Hand a timeout raised by the generator itself over unchanged."""
417+ exceptions : types .List [BaseException ] = []
418+
419+ async def generator () -> types .AsyncGenerator [int , None ]:
420+ """Yield one item and then fail with a timeout of its own."""
421+ yield 0
422+ raise asyncio .TimeoutError ('raised by the generator' )
423+
424+ def on_timeout (
425+ generator : types .AsyncGenerator [int , None ],
426+ timeout : types .Optional [types .delta_type ],
427+ total_timeout : types .Optional [types .delta_type ],
428+ exception : BaseException ,
429+ ) -> None :
430+ """Record the exception the detector reports."""
431+ exceptions .append (exception )
432+
433+ items : types .List [int ] = [
434+ i
435+ async for i in python_utils .aio_generator_timeout_detector (
436+ generator (), total_timeout = STALL , on_timeout = on_timeout
437+ )
438+ ]
439+
440+ assert items == [0 ]
441+ assert [str (exception ) for exception in exceptions ] == [
442+ 'raised by the generator'
443+ ]
444+
445+
446+ @pytest .mark .parametrize (
447+ 'timestamp,precision,expected' ,
448+ [
449+ (1 , datetime .timedelta (milliseconds = 100 ), '0:00:01' ),
450+ (60 , datetime .timedelta (milliseconds = 100 ), '0:01:00' ),
451+ (0.3 , datetime .timedelta (milliseconds = 100 ), '0:00:00.300000' ),
452+ (1.234 , datetime .timedelta (milliseconds = 1 ), '0:00:01.234000' ),
453+ (1.239 , datetime .timedelta (milliseconds = 10 ), '0:00:01.230000' ),
454+ (
455+ datetime .timedelta (seconds = 2 ),
456+ datetime .timedelta (milliseconds = 100 ),
457+ '0:00:02' ,
458+ ),
459+ (
460+ datetime .datetime (2000 , 1 , 2 , 3 , 4 , 5 ),
461+ datetime .timedelta (milliseconds = 100 ),
462+ '2000-01-02 03:04:05' ,
463+ ),
464+ (
465+ datetime .datetime (2000 , 1 , 2 , 3 , 4 , 5 , 678901 ),
466+ datetime .timedelta (milliseconds = 10 ),
467+ '2000-01-02 03:04:05.670000' ,
468+ ),
469+ ],
470+ )
471+ def test_format_time_sub_second_precision (
472+ timestamp : types .timestamp_type ,
473+ precision : datetime .timedelta ,
474+ expected : str ,
475+ ) -> None :
476+ """Keep a value that is on the precision grid where it is."""
477+ assert python_utils .format_time (timestamp , precision ) == expected
478+
479+
480+ @pytest .mark .parametrize (
481+ 'timestamp' ,
482+ [float ('nan' ), 'nan' , float ('inf' ), float ('-inf' ), 1e20 , 10 ** 30 ],
483+ )
484+ def test_format_time_placeholder_for_impossible_numbers (
485+ timestamp : types .timestamp_type ,
486+ ) -> None :
487+ """Print the placeholder for a number that is not a duration."""
488+ assert python_utils .format_time (timestamp ) == '--:--:--'
489+
490+
491+ def test_format_time_extreme_timedeltas () -> None :
492+ """Format the largest and the smallest timedelta without overflow."""
493+ largest : str = python_utils .format_time (datetime .timedelta .max )
494+ smallest : str = python_utils .format_time (datetime .timedelta .min )
495+
496+ assert largest == '999999999 days, 23:59:59'
497+ assert smallest == '-999999999 days, 0:00:00'
498+
499+
500+ @pytest .mark .parametrize (
501+ 'delta' ,
502+ [
503+ datetime .timedelta (microseconds = 1 ),
504+ datetime .timedelta (microseconds = - 1 ),
505+ datetime .timedelta (seconds = 437 , microseconds = 579262 ),
506+ datetime .timedelta (days = - 5 , microseconds = 1 ),
507+ datetime .timedelta (days = 1000000 , microseconds = 999999 ),
508+ datetime .timedelta .max ,
509+ ],
510+ )
511+ def test_timedelta_to_seconds_fraction_precision (
512+ delta : datetime .timedelta ,
513+ ) -> None :
514+ """Keep the precision of ``total_seconds`` for a fraction of a second."""
515+ seconds : types .Number = python_utils .timedelta_to_seconds (delta )
516+
517+ assert isinstance (seconds , float )
518+ assert seconds == delta .total_seconds ()
519+
520+
521+ @pytest .mark .parametrize (
522+ 'delta,expected' ,
523+ [
524+ (datetime .timedelta (0 ), 0 ),
525+ (datetime .timedelta (seconds = 1 ), 1 ),
526+ (datetime .timedelta (seconds = - 1 ), - 1 ),
527+ (datetime .timedelta (days = 1 ), 86400 ),
528+ (datetime .timedelta (days = 999999999 , seconds = 86399 ), 86399999999999 ),
529+ (datetime .timedelta .min , - 86399999913600 ),
530+ ],
531+ )
532+ def test_timedelta_to_seconds_whole_seconds_stay_int (
533+ delta : datetime .timedelta , expected : int
534+ ) -> None :
535+ """Return whole seconds as the exact ``int`` they have always been."""
536+ seconds : types .Number = python_utils .timedelta_to_seconds (delta )
537+
538+ assert isinstance (seconds , int )
539+ assert seconds == expected
540+
541+
542+ def test_timeout_generator_maximum_interval_first_sleep (
543+ fake_clock : clock .FakeClock ,
544+ ) -> None :
545+ """Hold the first sleep to ``maximum_interval`` as well."""
546+ items : types .List [str ] = list (
547+ python_utils .timeout_generator (
548+ timeout = 100 , interval = 10 , iterable = 'abc' , maximum_interval = 1
549+ )
550+ )
551+
552+ assert items == ['a' , 'b' , 'c' ]
553+ assert fake_clock .sleeps == [1 , 1 , 1 ]
554+
555+
556+ def test_timeout_generator_maximum_interval_zero (
557+ fake_clock : clock .FakeClock ,
558+ ) -> None :
559+ """Read a ``maximum_interval`` of zero as no maximum at all."""
560+ items : types .List [str ] = list (
561+ python_utils .timeout_generator (
562+ timeout = 100 ,
563+ interval = 1 ,
564+ iterable = 'abc' ,
565+ interval_multiplier = 2 ,
566+ maximum_interval = 0 ,
567+ )
568+ )
569+
570+ assert items == ['a' , 'b' , 'c' ]
571+ assert fake_clock .sleeps == [1 , 2 , 4 ]
572+
573+
574+ @pytest .mark .asyncio
575+ async def test_aio_timeout_generator_maximum_interval_first_sleep (
576+ monkeypatch : pytest .MonkeyPatch ,
577+ ) -> None :
578+ """Hold the first async sleep to ``maximum_interval`` as well."""
579+ sleeps : types .List [float ] = []
580+
581+ async def mock_sleep (delay : float ) -> None :
582+ """Record each requested delay instead of sleeping."""
583+ sleeps .append (delay )
584+
585+ monkeypatch .setattr (asyncio , 'sleep' , mock_sleep )
586+
587+ async def letters () -> types .AsyncGenerator [str , None ]:
588+ """Yield three letters without waiting."""
589+ letter : str
590+ for letter in 'abc' :
591+ yield letter
592+
593+ items : types .List [str ] = [
594+ item
595+ async for item in python_utils .aio_timeout_generator (
596+ timeout = 100 , interval = 10 , iterable = letters , maximum_interval = 1
597+ )
598+ ]
599+
600+ assert items == ['a' , 'b' , 'c' ]
601+ assert sleeps == [1 , 1 , 1 ]
0 commit comments