[tests] make test code generic over type of weighted stream

We're going to add a variant of this in the next commit.

Author: sunshowers

src/future_queue.rs

@@ -6,7 +6,6 @@ use futures_util::{
     Future, Stream, StreamExt as _,
 };
 use pin_project_lite::pin_project;
-use private::WeightedFuture;
 use std::{
     fmt,
     pin::Pin,
@@ -167,25 +166,31 @@ where
     }
 }
 
-mod private {
-    use futures_util::Future;
+/// A trait for types which can be converted into a `Future` and a weight.
+///
+/// Provided in case it's necessary. This trait is only implemented for `(usize, impl Future)`.
+pub trait WeightedFuture: private::Sealed {
+    type Future: Future;
 
-    pub trait WeightedFuture {
-        type Future: Future;
+    /// Turns this trait into its components.
+    fn into_components(self) -> (usize, Self::Future);
+}
 
-        fn into_components(self) -> (usize, Self::Future);
-    }
+mod private {
+    pub trait Sealed {}
+}
 
-    impl<Fut> WeightedFuture for (usize, Fut)
-    where
-        Fut: Future,
-    {
-        type Future = Fut;
+impl<Fut> private::Sealed for (usize, Fut) where Fut: Future {}
 
-        #[inline]
-        fn into_components(self) -> (usize, Self::Future) {
-            self
-        }
+impl<Fut> WeightedFuture for (usize, Fut)
+where
+    Fut: Future,
+{
+    type Future = Fut;
+
+    #[inline]
+    fn into_components(self) -> (usize, Self::Future) {
+        self
     }
 }
 

src/lib.rs

@@ -98,7 +98,15 @@
 
 mod future_queue;
 
-pub use future_queue::FutureQueue;
+pub use crate::future_queue::FutureQueue;
+
+/// Traits to aid in type definitions.
+///
+/// These traits are normally not required by end-user code, but may be necessary for some generic
+/// code.
+pub mod traits {
+    pub use crate::future_queue::WeightedFuture;
+}
 
 use futures_util::{Future, Stream};
 

tests/test_properties.rs

@@ -1,36 +1,127 @@
 // Copyright (c) The future-queue Contributors
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-use future_queue::StreamExt as _;
-use futures::{stream, StreamExt as _};
+use future_queue::{traits::WeightedFuture, FutureQueue, StreamExt as _};
+use futures::{future::BoxFuture, stream, Future, FutureExt, Stream, StreamExt as _};
 use proptest::prelude::*;
 use proptest_derive::Arbitrary;
-use std::time::Duration;
+use std::{pin::Pin, time::Duration};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
 #[derive(Clone, Debug, Arbitrary)]
-struct TestState {
+struct TestState<G: GroupSpec> {
     #[proptest(strategy = "1usize..64")]
     max_weight: usize,
     #[proptest(strategy = "prop::collection::vec(TestFutureDesc::arbitrary(), 0..512usize)")]
-    future_descriptions: Vec<TestFutureDesc>,
+    future_descriptions: Vec<TestFutureDesc<G>>,
 }
 
 #[derive(Copy, Clone, Debug, Arbitrary)]
-struct TestFutureDesc {
+struct TestFutureDesc<G: GroupSpec> {
     #[proptest(strategy = "duration_strategy()")]
     start_delay: Duration,
     #[proptest(strategy = "duration_strategy()")]
     delay: Duration,
     #[proptest(strategy = "0usize..8")]
     weight: usize,
+    #[allow(dead_code)]
+    group: G,
 }
 
 fn duration_strategy() -> BoxedStrategy<Duration> {
     // Allow for a delay between 0ms and 1000ms uniformly at random.
     (0u64..1000).prop_map(Duration::from_millis).boxed()
 }
 
+trait GroupSpec: Arbitrary + Send + Copy + 'static {
+    type Item: Send;
+    type CheckState: Default;
+
+    fn create_stream<St>(stream: St, state: &TestState<Self>) -> BoxedWeightedStream<()>
+    where
+        St: Stream<Item = Self::Item> + Send + 'static;
+
+    fn create_stream_item(
+        desc: &TestFutureDesc<Self>,
+        future: impl Future<Output = ()> + Send + 'static,
+    ) -> Self::Item;
+
+    fn check_started(
+        check_state: &mut Self::CheckState,
+        desc: &TestFutureDesc<Self>,
+        state: &TestState<Self>,
+    );
+
+    fn check_finished(
+        check_state: &mut Self::CheckState,
+        desc: &TestFutureDesc<Self>,
+        state: &TestState<Self>,
+    );
+}
+
+trait WeightedStream: Stream {
+    fn current_weight(&self) -> usize;
+}
+
+impl<St, Fut> WeightedStream for FutureQueue<St>
+where
+    St: Stream<Item = Fut>,
+    Fut: WeightedFuture,
+{
+    fn current_weight(&self) -> usize {
+        self.current_weight()
+    }
+}
+
+type BoxedWeightedStream<Item> = Pin<Box<dyn WeightedStream<Item = Item> + Send>>;
+
+impl GroupSpec for () {
+    type Item = (usize, BoxFuture<'static, ()>);
+    type CheckState = NonGroupedCheckState;
+
+    fn create_stream<St>(stream: St, state: &TestState<Self>) -> BoxedWeightedStream<()>
+    where
+        St: Stream<Item = Self::Item> + Send + 'static,
+    {
+        Box::pin(stream.future_queue(state.max_weight))
+    }
+
+    fn create_stream_item(
+        desc: &TestFutureDesc<Self>,
+        future: impl Future<Output = ()> + Send + 'static,
+    ) -> Self::Item {
+        (desc.weight, future.boxed())
+    }
+
+    fn check_started(
+        check_state: &mut Self::CheckState,
+        desc: &TestFutureDesc<Self>,
+        state: &TestState<Self>,
+    ) {
+        // Check that current_weight doesn't go over the limit.
+        assert!(
+            check_state.current_weight < state.max_weight,
+            "current weight {} exceeds max weight {}",
+            check_state.current_weight,
+            state.max_weight,
+        );
+        check_state.current_weight += desc.weight;
+    }
+
+    fn check_finished(
+        check_state: &mut Self::CheckState,
+        desc: &TestFutureDesc<Self>,
+        _state: &TestState<Self>,
+    ) {
+        check_state.current_weight -= desc.weight;
+    }
+}
+
+#[derive(Debug, Default)]
+struct NonGroupedCheckState {
+    current_weight: usize,
+}
+
 #[test]
 fn test_examples() {
     let state = TestState {
@@ -39,25 +130,26 @@ fn test_examples() {
             start_delay: Duration::ZERO,
             delay: Duration::ZERO,
             weight: 0,
+            group: (),
         }],
     };
     test_future_queue_impl(state);
 }
 
 proptest! {
     #[test]
-    fn proptest_future_queue(state: TestState) {
+    fn proptest_future_queue(state: TestState<()>) {
         test_future_queue_impl(state)
     }
 }
 
 #[derive(Clone, Copy, Debug)]
-enum FutureEvent {
-    Started(usize, TestFutureDesc),
-    Finished(usize, TestFutureDesc),
+enum FutureEvent<G: GroupSpec> {
+    Started(usize, TestFutureDesc<G>),
+    Finished(usize, TestFutureDesc<G>),
 }
 
-fn test_future_queue_impl(state: TestState) {
+fn test_future_queue_impl<G: GroupSpec>(state: TestState<G>) {
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_time()
         .start_paused(true)
@@ -88,7 +180,7 @@ fn test_future_queue_impl(state: TestState) {
                         .expect("receiver held open by loop");
                 };
                 // Errors should never occur here.
-                if let Err(err) = future_sender.send((desc.weight, delay_fut)) {
+                if let Err(err) = future_sender.send(G::create_stream_item(&desc, delay_fut)) {
                     panic!("future_receiver held open by loop: {}", err);
                 }
             }
@@ -102,11 +194,11 @@ fn test_future_queue_impl(state: TestState) {
 
     let mut completed_map = vec![false; state.future_descriptions.len()];
     let mut last_started_id: Option<usize> = None;
-    let mut current_weight = 0;
+    let mut check_state = G::CheckState::default();
 
     runtime.block_on(async move {
         // Record values that have been completed in this map.
-        let mut stream = stream.future_queue(state.max_weight);
+        let mut stream = G::create_stream(stream, &state);
         let mut receiver_done = false;
         loop {
             tokio::select! {
@@ -122,19 +214,12 @@ fn test_future_queue_impl(state: TestState) {
                             assert_eq!(expected_id, id, "expected future id to start != actual id that started");
                             last_started_id = Some(id);
 
-                            // Check that the current weight doesn't go over the limit.
-                            assert!(
-                                current_weight < state.max_weight,
-                                "current weight {} exceeds max weight {}",
-                                current_weight,
-                                state.max_weight,
-                            );
-                            current_weight += desc.weight;
+                            G::check_started(&mut check_state, &desc, &state);
                         }
                         Some(FutureEvent::Finished(id, desc)) => {
                             // Record that this value was completed.
                             completed_map[id] = true;
-                            current_weight -= desc.weight;
+                            G::check_finished(&mut check_state, &desc, &state);
                         }
                         None => {
                             // All futures finished -- going to check for completion in stream.next() below.