test: add some tests of recent stop & drop additions

Author: obmarg

tests/cynic-tests.rs

@@ -1,7 +1,9 @@
 use std::{future::IntoFuture, time::Duration};
 
 use assert_matches::assert_matches;
+use cynic::StreamingOperation;
 use futures_lite::{StreamExt, future};
+use graphql_ws_client::Subscription;
 use subscription_server::SubscriptionServer;
 use tokio::time::sleep;
 
@@ -54,29 +56,14 @@ struct BooksChangedSubscription {
 
 #[tokio::test]
 async fn main_test() {
-    use async_tungstenite::tungstenite::{client::IntoClientRequest, http::HeaderValue};
-
     let server = SubscriptionServer::start().await;
 
-    sleep(Duration::from_millis(20)).await;
-
-    let mut request = server.websocket_url().into_client_request().unwrap();
-    request.headers_mut().insert(
-        "Sec-WebSocket-Protocol",
-        HeaderValue::from_str("graphql-transport-ws").unwrap(),
-    );
-
-    let (connection, _) = async_tungstenite::tokio::connect_async(request)
-        .await
-        .unwrap();
-
-    println!("Connected");
-
-    let (client, actor) = graphql_ws_client::Client::build(connection).await.unwrap();
+    let client_builder = server.client_builder().await;
+    let (client, actor) = client_builder.await.unwrap();
 
     tokio::spawn(actor.into_future());
 
-    let stream = client.subscribe(build_query()).await.unwrap();
+    let mut stream = client.subscribe(build_query()).await.unwrap();
 
     sleep(Duration::from_millis(100)).await;
 
@@ -95,50 +82,16 @@ async fn main_test() {
         },
     ];
 
-    future::zip(
-        async {
-            for update in &updates {
-                server.send(update.to_owned()).unwrap();
-            }
-        },
-        async {
-            let received_updates = stream.take(updates.len()).collect::<Vec<_>>().await;
-
-            for (expected, update) in updates.iter().zip(received_updates) {
-                let update = update.unwrap();
-                assert_matches!(update.errors, None);
-                let data = update.data.unwrap();
-                assert_eq!(data.books.id.inner(), expected.id.0);
-            }
-        },
-    )
-    .await;
+    send_and_verify_updates(&server, &updates, &mut stream).await;
 }
 
 #[tokio::test]
 async fn oneshot_operation_test() {
-    use async_tungstenite::tungstenite::{client::IntoClientRequest, http::HeaderValue};
-
     let server = SubscriptionServer::start().await;
 
-    sleep(Duration::from_millis(20)).await;
-
-    let mut request = server.websocket_url().into_client_request().unwrap();
-    request.headers_mut().insert(
-        "Sec-WebSocket-Protocol",
-        HeaderValue::from_str("graphql-transport-ws").unwrap(),
-    );
-
-    let (connection, _) = async_tungstenite::tokio::connect_async(request)
-        .await
-        .unwrap();
+    let client_builder = server.client_builder().await;
 
-    println!("Connected");
-
-    let stream = graphql_ws_client::Client::build(connection)
-        .subscribe(build_query())
-        .await
-        .unwrap();
+    let mut stream = client_builder.subscribe(build_query()).await.unwrap();
 
     let updates = [
         subscription_server::BookChanged {
@@ -155,15 +108,115 @@ async fn oneshot_operation_test() {
         },
     ];
 
+    send_and_verify_updates(&server, &updates, &mut stream).await;
+}
+
+#[tokio::test]
+async fn test_client_stop() {
+    let server = SubscriptionServer::start().await;
+
+    let client_builder = server.client_builder().await;
+    let (client, actor) = client_builder.await.unwrap();
+
+    tokio::spawn(actor.into_future());
+
+    let mut stream = client.subscribe(build_query()).await.unwrap();
+
+    sleep(Duration::from_millis(10)).await;
+
+    let updates = [subscription_server::BookChanged {
+        id: "123".into(),
+        book: None,
+    }];
+
+    send_and_verify_updates(&server, &updates, &mut stream).await;
+
+    assert_eq!(server.subscriber_count(), 1);
+
+    client.stop(stream.id()).await.unwrap();
+
+    sleep(Duration::from_millis(10)).await;
+
+    assert_eq!(server.subscriber_count(), 0);
+
+    assert!(stream.next().await.is_none());
+}
+
+#[tokio::test]
+async fn test_subscription_stop() {
+    let server = SubscriptionServer::start().await;
+
+    let client_builder = server.client_builder().await;
+    let (client, actor) = client_builder.await.unwrap();
+
+    tokio::spawn(actor.into_future());
+
+    let mut stream = client.subscribe(build_query()).await.unwrap();
+
+    sleep(Duration::from_millis(10)).await;
+
+    let updates = [subscription_server::BookChanged {
+        id: "123".into(),
+        book: None,
+    }];
+
+    send_and_verify_updates(&server, &updates, &mut stream).await;
+
+    assert_eq!(server.subscriber_count(), 1);
+
+    stream.stop().await.unwrap();
+
+    sleep(Duration::from_millis(10)).await;
+
+    assert_eq!(server.subscriber_count(), 0);
+}
+
+#[tokio::test]
+async fn test_subscription_stops_on_drop() {
+    let server = SubscriptionServer::start().await;
+
+    let client_builder = server.client_builder().await;
+    let (client, actor) = client_builder.await.unwrap();
+
+    tokio::spawn(actor.into_future());
+
+    let mut stream = client.subscribe(build_query()).await.unwrap();
+
+    sleep(Duration::from_millis(10)).await;
+
+    let updates = [subscription_server::BookChanged {
+        id: "123".into(),
+        book: None,
+    }];
+
+    send_and_verify_updates(&server, &updates, &mut stream).await;
+
+    assert_eq!(server.subscriber_count(), 1);
+
+    drop(stream);
+
+    sleep(Duration::from_millis(10)).await;
+
+    assert_eq!(server.subscriber_count(), 0);
+}
+
+async fn send_and_verify_updates(
+    server: &SubscriptionServer,
+    updates: &[subscription_server::BookChanged],
+    stream: &mut Subscription<StreamingOperation<BooksChangedSubscription, BooksChangedVariables>>,
+) {
     future::zip(
         async {
             sleep(Duration::from_millis(10)).await;
-            for update in &updates {
+            for update in updates {
                 server.send(update.to_owned()).unwrap();
             }
         },
         async {
-            let received_updates = stream.take(updates.len()).collect::<Vec<_>>().await;
+            let mut received_updates = Vec::new();
+            for _ in updates {
+                received_updates.push(stream.next().await.unwrap());
+            }
 
             for (expected, update) in updates.iter().zip(received_updates) {
                 let update = update.unwrap();

tests/subscription_server/mod.rs

@@ -1,19 +1,30 @@
 // TODO: Start a server w/ async-graphql.
 // Query that server...
 
+use std::{
+    pin::Pin,
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+    time::Duration,
+};
+
 use async_graphql::{EmptyMutation, ID, Object, Schema, SimpleObject, Subscription};
 use async_graphql_axum::{GraphQLRequest, GraphQLResponse, GraphQLSubscription};
-use axum::{Router, extract::Extension, routing::post};
+use axum::{Router, extract::Extension, http::HeaderValue, routing::post};
 use futures_lite::{Stream, StreamExt};
 use tokio::sync::broadcast::Sender;
 use tokio_stream::wrappers::BroadcastStream;
+use tungstenite::client::IntoClientRequest;
 
 pub type BooksSchema = Schema<QueryRoot, EmptyMutation, SubscriptionRoot>;
 
 pub struct SubscriptionServer {
     shutdown: Option<tokio::sync::oneshot::Sender<()>>,
     port: u16,
     sender: Sender<BookChanged>,
+    subscriber_count: Arc<AtomicUsize>,
 }
 
 impl Drop for SubscriptionServer {
@@ -27,12 +38,14 @@ impl Drop for SubscriptionServer {
 impl SubscriptionServer {
     pub async fn start() -> SubscriptionServer {
         let (channel, _) = tokio::sync::broadcast::channel(16);
+        let subscriber_count = Arc::new(AtomicUsize::new(0));
 
         let schema = Schema::build(
             QueryRoot,
             EmptyMutation,
             SubscriptionRoot {
                 channel: channel.clone(),
+                subscriber_count: Arc::clone(&subscriber_count),
             },
         )
         .finish();
@@ -56,23 +69,48 @@ impl SubscriptionServer {
                 .unwrap();
         });
 
+        tokio::time::sleep(Duration::from_millis(20)).await;
+
         SubscriptionServer {
             port,
             shutdown: Some(shutdown_sender),
             sender: channel,
+            subscriber_count,
         }
     }
 
     pub fn websocket_url(&self) -> String {
         format!("ws://localhost:{}/ws", self.port)
     }
 
+    #[allow(unused)]
+    pub fn subscriber_count(&self) -> usize {
+        self.subscriber_count.load(Ordering::Relaxed)
+    }
+
     pub fn send(
         &self,
         change: BookChanged,
     ) -> Result<(), tokio::sync::broadcast::error::SendError<BookChanged>> {
         self.sender.send(change).map(|_| ())
     }
+
+    #[allow(unused)]
+    pub async fn client_builder(&self) -> graphql_ws_client::ClientBuilder {
+        let mut request = self.websocket_url().into_client_request().unwrap();
+        request.headers_mut().insert(
+            "Sec-WebSocket-Protocol",
+            HeaderValue::from_str("graphql-transport-ws").unwrap(),
+        );
+
+        let (connection, _) = async_tungstenite::tokio::connect_async(request)
+            .await
+            .unwrap();
+
+        println!("Connected");
+
+        graphql_ws_client::Client::build(connection)
+    }
 }
 
 #[axum_macros::debug_handler]
@@ -110,12 +148,45 @@ enum MutationType {
 
 pub struct SubscriptionRoot {
     channel: Sender<BookChanged>,
+    subscriber_count: Arc<AtomicUsize>,
 }
 
 #[Subscription]
 impl SubscriptionRoot {
     async fn books(&self, _mutation_type: MutationType) -> impl Stream<Item = BookChanged> {
         println!("Subscription received");
-        BroadcastStream::new(self.channel.subscribe()).filter_map(Result::ok)
+        self.subscriber_count.fetch_add(1, Ordering::Relaxed);
+        TrackedBroadcastStream {
+            inner: BroadcastStream::new(self.channel.subscribe()).filter_map(Result::ok),
+            count: Arc::clone(&self.subscriber_count),
+        }
+    }
+}
+
+#[pin_project::pin_project(PinnedDrop)]
+pub struct TrackedBroadcastStream<T> {
+    #[pin]
+    inner: T,
+    count: Arc<AtomicUsize>,
+}
+
+impl<T> Stream for TrackedBroadcastStream<T>
+where
+    T: Stream,
+{
+    type Item = <T as Stream>::Item;
+
+    fn poll_next(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.project().inner.poll_next(cx)
+    }
+}
+
+#[pin_project::pinned_drop]
+impl<T> PinnedDrop for TrackedBroadcastStream<T> {
+    fn drop(self: Pin<&mut Self>) {
+        self.project().count.fetch_sub(1, Ordering::Relaxed);
     }
 }