#[ cfg( all( test, feature = "streaming-control" ) ) ]
mod tests
{
use api_huggingface::
{
streaming_control::wrap_stream,
error::{ Result, HuggingFaceError },
};
use tokio::sync::mpsc;
use futures_util::StreamExt;
fn create_mock_stream() -> mpsc::Receiver< Result< String > >
{
let ( tx, rx ) = mpsc::channel( 10 );
tokio::spawn( async move
{
for i in 0..10
{
if tx.send( Ok( format!( "chunk_{i}" ) ) ).await.is_err()
{
break;
}
tokio::time::sleep( tokio::time::Duration::from_millis( 50 ) ).await;
}
});
rx
}
#[ tokio::test ]
async fn test_streaming_no_control()
{
let receiver = create_mock_stream();
let ( mut stream, _control ) = wrap_stream( receiver );
let mut count = 0;
while let Some( result ) = stream.next().await
{
assert!( result.is_ok() );
count += 1;
}
assert_eq!( count, 10 );
}
#[ tokio::test ]
async fn test_pause_resume()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
for _ in 0..2
{
assert!( stream.next().await.is_some() );
}
control.pause().await.unwrap();
tokio::time::sleep( tokio::time::Duration::from_millis( 200 ) ).await;
control.resume().await.unwrap();
let mut count = 2;
while let Some( result ) = stream.next().await
{
assert!( result.is_ok() );
count += 1;
}
assert_eq!( count, 10 );
}
#[ tokio::test ]
async fn test_cancel()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
for _ in 0..3
{
assert!( stream.next().await.is_some() );
}
control.cancel().await.unwrap();
assert!( stream.next().await.is_none() );
}
#[ tokio::test ]
async fn test_is_paused()
{
let receiver = create_mock_stream();
let ( stream, control ) = wrap_stream( receiver );
assert!( !stream.is_paused() );
assert!( stream.is_running() );
control.pause().await.unwrap();
control.cancel().await.unwrap();
}
#[ tokio::test ]
async fn test_is_cancelled()
{
let receiver = create_mock_stream();
let ( stream, control ) = wrap_stream( receiver );
assert!( !stream.is_cancelled() );
assert!( stream.is_running() );
control.cancel().await.unwrap();
}
#[ tokio::test ]
async fn test_try_pause()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
assert!( control.try_pause().is_ok() );
control.cancel().await.unwrap();
while stream.next().await.is_some() {}
}
#[ tokio::test ]
async fn test_try_resume()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
control.pause().await.unwrap();
assert!( control.try_resume().is_ok() );
control.cancel().await.unwrap();
while stream.next().await.is_some() {}
}
#[ tokio::test ]
async fn test_try_cancel()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
assert!( control.try_cancel().is_ok() );
assert!( stream.next().await.is_none() );
}
#[ tokio::test ]
async fn test_multiple_pause_resume()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
for cycle in 0..3
{
assert!( stream.next().await.is_some(), "Failed at cycle {cycle}" );
control.pause().await.unwrap();
tokio::time::sleep( tokio::time::Duration::from_millis( 50 ) ).await;
control.resume().await.unwrap();
}
control.cancel().await.unwrap();
assert!( stream.next().await.is_none() );
}
#[ tokio::test ]
async fn test_control_after_natural_end()
{
let ( tx, rx ) = mpsc::channel( 10 );
tokio::spawn( async move
{
tx.send( Ok( "chunk_0".to_string() ) ).await.ok();
tx.send( Ok( "chunk_1".to_string() ) ).await.ok();
});
let ( mut stream, _control ) = wrap_stream( rx );
while stream.next().await.is_some() {}
assert!( stream.next().await.is_none() );
}
#[ tokio::test ]
async fn test_control_handle_clone()
{
let receiver = create_mock_stream();
let ( mut stream, control ) = wrap_stream( receiver );
let control_clone = control.clone();
control.pause().await.unwrap();
control_clone.resume().await.unwrap();
assert!( stream.next().await.is_some() );
control.cancel().await.unwrap();
while stream.next().await.is_some() {}
}
#[ tokio::test ]
async fn test_streaming_with_errors()
{
let ( tx, rx ) = mpsc::channel( 10 );
tokio::spawn( async move
{
tx.send( Ok( "chunk_0".to_string() ) ).await.ok();
tx.send( Err( HuggingFaceError::Stream( "test error".to_string() ) ) ).await.ok();
tx.send( Ok( "chunk_2".to_string() ) ).await.ok();
});
let ( mut stream, control ) = wrap_stream( rx );
assert!( stream.next().await.unwrap().is_ok() );
assert!( stream.next().await.unwrap().is_err() );
assert!( stream.next().await.unwrap().is_ok() );
control.cancel().await.unwrap();
}
#[ tokio::test ]
async fn test_pause_during_errors()
{
let ( tx, rx ) = mpsc::channel( 10 );
tokio::spawn( async move
{
for i in 0..5
{
if i == 2
{
tx.send( Err( HuggingFaceError::Stream( "error".to_string() ) ) ).await.ok();
}
else
{
tx.send( Ok( format!( "chunk_{i}" ) ) ).await.ok();
}
tokio::time::sleep( tokio::time::Duration::from_millis( 50 ) ).await;
}
});
let ( mut stream, control ) = wrap_stream( rx );
assert!( stream.next().await.unwrap().is_ok() );
control.pause().await.unwrap();
tokio::time::sleep( tokio::time::Duration::from_millis( 100 ) ).await;
control.resume().await.unwrap();
assert!( stream.next().await.unwrap().is_ok() );
assert!( stream.next().await.unwrap().is_err() );
control.cancel().await.unwrap();
}
#[ test ]
fn test_debug_implementations()
{
let ( _tx, rx ) = mpsc::channel::< Result< String > >( 10 );
let ( stream, handle ) = wrap_stream( rx );
let _ = format!( "{stream:?}" );
let _ = format!( "{handle:?}" );
}
}