@@ -396,6 +396,170 @@ describe('streamedQuery', () => {
396396 } )
397397 } )
398398
399+ it . each ( [ 'append' , 'replace' ] as const ) (
400+ 'should stop %s writes after explicit cancellation when the stream ignores the signal' ,
401+ async ( refetchMode ) => {
402+ const key = createKey ( )
403+ queryClient . setQueryData ( key , [ 99 ] )
404+ const observer = new QueryObserver ( queryClient , {
405+ queryKey : key ,
406+ queryFn : streamedQuery ( {
407+ streamFn : ( ) => createAsyncNumberGenerator ( 3 ) ,
408+ refetchMode,
409+ } ) ,
410+ persister : ( queryFn , context ) => queryFn ( { ...context } ) ,
411+ } )
412+
413+ const unsubscribe = observer . subscribe ( vi . fn ( ) )
414+ await vi . advanceTimersByTimeAsync ( 50 )
415+ await queryClient . cancelQueries ( { queryKey : key } )
416+ const cancelledData = queryClient . getQueryData ( key )
417+ expect ( cancelledData ) . toStrictEqual (
418+ refetchMode === 'append' ? [ 99 , 0 ] : [ 99 ] ,
419+ )
420+
421+ await vi . advanceTimersByTimeAsync ( 150 )
422+ expect ( queryClient . getQueryData ( key ) ) . toStrictEqual ( cancelledData )
423+ unsubscribe ( )
424+ } ,
425+ )
426+
427+ it ( 'should preserve a non-streamed replacement when the old stream ignores the signal' , async ( ) => {
428+ const key = createKey ( )
429+ const queryFn = streamedQuery ( {
430+ streamFn : ( ) => createAsyncNumberGenerator ( 3 ) ,
431+ } )
432+ const observer = new QueryObserver ( queryClient , {
433+ queryKey : key ,
434+ queryFn : ( context ) => queryFn ( context ) ,
435+ } )
436+
437+ const unsubscribe = observer . subscribe ( vi . fn ( ) )
438+ await vi . advanceTimersByTimeAsync ( 50 )
439+ observer . setOptions ( { queryKey : key , queryFn : ( ) => [ 100 ] } )
440+ await observer . refetch ( )
441+ await vi . advanceTimersByTimeAsync ( 150 )
442+
443+ expect ( queryClient . getQueryData ( key ) ) . toStrictEqual ( [ 100 ] )
444+ unsubscribe ( )
445+ } )
446+
447+ it ( 'should preserve the new stream when the old stream ignores the signal' , async ( ) => {
448+ const key = createKey ( )
449+ let offset = 0
450+ const observer = new QueryObserver ( queryClient , {
451+ queryKey : key ,
452+ queryFn : streamedQuery ( {
453+ streamFn : ( ) => {
454+ const stream = createAsyncNumberGenerator ( 3 , offset )
455+ offset += 100
456+ return stream
457+ } ,
458+ } ) ,
459+ } )
460+
461+ const unsubscribe = observer . subscribe ( vi . fn ( ) )
462+ await vi . advanceTimersByTimeAsync ( 50 )
463+ void observer . refetch ( )
464+ await vi . advanceTimersByTimeAsync ( 150 )
465+
466+ expect ( queryClient . getQueryData ( key ) ) . toStrictEqual ( [ 100 , 101 , 102 ] )
467+ unsubscribe ( )
468+ } )
469+
470+ it ( 'should ignore a chunk queued just before cancellation' , async ( ) => {
471+ const key = createKey ( )
472+ queryClient . setQueryData ( key , [ 99 ] )
473+ let resolveChunk : ( chunk : IteratorResult < number > ) => void = ( ) => {
474+ throw new Error ( 'The iterator has not started' )
475+ }
476+ const next = vi . fn < ( ) => Promise < IteratorResult < number > > > (
477+ ( ) =>
478+ new Promise < IteratorResult < number > > ( ( resolve ) => {
479+ resolveChunk = resolve
480+ } ) ,
481+ )
482+ const observer = new QueryObserver ( queryClient , {
483+ queryKey : key ,
484+ queryFn : streamedQuery ( {
485+ refetchMode : 'append' ,
486+ streamFn : ( ) => ( {
487+ [ Symbol . asyncIterator ] : ( ) => ( { next } ) ,
488+ } ) ,
489+ } ) ,
490+ } )
491+
492+ const unsubscribe = observer . subscribe ( vi . fn ( ) )
493+ await vi . advanceTimersByTimeAsync ( 0 )
494+ resolveChunk ( { done : false , value : 1 } )
495+ void queryClient . cancelQueries ( { queryKey : key } )
496+ await vi . advanceTimersByTimeAsync ( 0 )
497+
498+ expect ( queryClient . getQueryData ( key ) ) . toStrictEqual ( [ 99 ] )
499+ expect ( next ) . toHaveBeenCalledTimes ( 1 )
500+ unsubscribe ( )
501+ } )
502+
503+ it . each ( [ 'queryFn' , 'persister' ] as const ) (
504+ 'should not start a stream after cancellation during a delayed %s' ,
505+ async ( delayIn ) => {
506+ const key = createKey ( )
507+ queryClient . setQueryData ( key , [ 99 ] )
508+ const streamFn = vi . fn < ( ) => AsyncIterable < number > > ( ( ) =>
509+ createAsyncNumberGenerator ( 1 ) ,
510+ )
511+ const queryFn = streamedQuery ( { streamFn } )
512+ const observer = new QueryObserver ( queryClient , {
513+ queryKey : key ,
514+ queryFn :
515+ delayIn === 'queryFn'
516+ ? async ( context ) => {
517+ await sleep ( 50 )
518+ return queryFn ( context )
519+ }
520+ : queryFn ,
521+ persister :
522+ delayIn === 'persister'
523+ ? async ( fn , context ) => {
524+ await sleep ( 50 )
525+ return fn ( { ...context } )
526+ }
527+ : undefined ,
528+ } )
529+
530+ const unsubscribe = observer . subscribe ( vi . fn ( ) )
531+ await queryClient . cancelQueries ( { queryKey : key } )
532+ queryClient . setQueryData ( key , [ 100 ] )
533+ await vi . advanceTimersByTimeAsync ( 100 )
534+
535+ expect ( streamFn ) . not . toHaveBeenCalled ( )
536+ expect ( queryClient . getQueryData ( key ) ) . toStrictEqual ( [ 100 ] )
537+ unsubscribe ( )
538+ } ,
539+ )
540+
541+ it ( 'should keep a custom signal getter lazy' , async ( ) => {
542+ const key = createKey ( )
543+ const controller = new AbortController ( )
544+ const getSignal = vi . fn < ( ) => AbortSignal > ( ( ) => controller . signal )
545+ const queryFn = streamedQuery ( {
546+ streamFn : ( ) => createAsyncNumberGenerator ( 1 ) ,
547+ } )
548+
549+ const result = queryFn ( {
550+ client : queryClient ,
551+ queryKey : key ,
552+ meta : undefined ,
553+ get signal ( ) {
554+ return getSignal ( )
555+ } ,
556+ } )
557+ await vi . advanceTimersByTimeAsync ( 50 )
558+
559+ await expect ( result ) . resolves . toStrictEqual ( [ 0 ] )
560+ expect ( getSignal ) . not . toHaveBeenCalled ( )
561+ } )
562+
399563 it ( 'should support custom reducer' , async ( ) => {
400564 const key = createKey ( )
401565
0 commit comments