diff --git a/Cargo.lock b/Cargo.lock index 8f065bd..7edcbd1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -220,6 +220,28 @@ dependencies = [ "event-listener", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "async-trait" version = "0.1.91" @@ -282,7 +304,7 @@ dependencies = [ "sha1", "sync_wrapper", "tokio", - "tokio-tungstenite", + "tokio-tungstenite 0.29.0", "tower", "tower-layer", "tower-service", @@ -484,6 +506,22 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -889,15 +927,20 @@ version = "0.0.0" dependencies = [ "async-graphql", "async-graphql-axum", + "async-stream", "axum", "clap", "config", + "futures-util", "http-body-util", "httpmock", "reqwest", + "rmp-serde", + "rstest", "serde", "serde_json", "tokio", + "tokio-tungstenite 0.28.0", "tower", "tracing", "tracing-subscriber", @@ -905,6 +948,12 @@ dependencies = [ "uuid", ] +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + [[package]] name = "h2" version = "0.4.15" @@ -1443,6 +1492,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + [[package]] name = "ordered-multimap" version = "0.7.3" @@ -1771,6 +1826,12 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + [[package]] name = "reqwest" version = "0.12.28" @@ -1826,6 +1887,25 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmp" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" +dependencies = [ + "rmp", + "serde", +] + [[package]] name = "ron" version = "0.12.2" @@ -1840,6 +1920,35 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.119", + "unicode-ident", +] + [[package]] name = "rust-ini" version = "0.21.3" @@ -1856,6 +1965,15 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "1.1.4" @@ -1883,6 +2001,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + [[package]] name = "rustls-pki-types" version = "1.15.1" @@ -1916,12 +2046,50 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.229" @@ -2317,6 +2485,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite 0.28.0", +] + [[package]] name = "tokio-tungstenite" version = "0.29.0" @@ -2326,7 +2510,7 @@ dependencies = [ "futures-util", "log", "tokio", - "tungstenite", + "tungstenite 0.29.0", ] [[package]] @@ -2497,6 +2681,25 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.5", + "rustls", + "rustls-pki-types", + "sha1", + "thiserror", + "utf-8", +] + [[package]] name = "tungstenite" version = "0.29.0" @@ -2568,6 +2771,12 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8_iter" version = "1.0.4" diff --git a/Cargo.toml b/Cargo.toml index 6ee276d..f088f71 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,8 +20,13 @@ clap = { version = "4.5.48", features = ["derive"] } uuid = { version = "1.18.1", features = ["serde"] } tracing = "0.1.41" tracing-subscriber = "0.3.20" +tokio-tungstenite = { version = "0.28.0", features = ["rustls-tls-native-roots"] } +rmp-serde = "1.3.1" +futures-util = "0.3.32" +async-stream = "0.3.6" [dev-dependencies] http-body-util = "0.1.3" httpmock = "0.8.2" +rstest = "0.26.1" tower = "0.5.2" diff --git a/src/clients.rs b/src/clients.rs index e0cecd5..61a81b3 100644 --- a/src/clients.rs +++ b/src/clients.rs @@ -1,13 +1,18 @@ use std::borrow::Cow; use std::fmt; +use futures_util::{Stream, StreamExt}; #[cfg(test)] use httpmock::MockServer; use reqwest::header::HeaderMap; use reqwest::{Client, Url}; use serde::de::DeserializeOwned; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::{debug, info, instrument}; +use crate::model::subscription::TiledEvent; use crate::model::{app, node, table}; pub type ClientResult = Result; @@ -30,6 +35,7 @@ impl TiledClient { address, } } + #[instrument(skip(self, headers))] async fn request( &self, @@ -130,6 +136,43 @@ impl TiledClient { .await } + pub(crate) fn stream_events( + &self, + node: Option, + headers: Option, + ) -> impl Stream> { + let scheme = match self.address.scheme() { + "https" => "wss", + _ => "ws", + }; + + let path = format!("api/v1/stream/single/{}", node.as_deref().unwrap_or("")); + + let mut url = self.address.join(&path).expect("Invalid stream path"); + url.set_scheme(scheme) + .expect("ws and wss are valid schemes"); + url.set_query(Some("envelope_format=msgpack")); + + let mut request = url + .as_str() + .into_client_request() + .expect("valid URLs can be converted to URIs needed for Request"); + if let Some(headers) = headers { + request.headers_mut().extend(headers); + } + async_stream::try_stream! { + info!("Connecting to WebSocket: {}", url); + let (ws,_) = connect_async(request).await.map_err(|e| SubscriptionError::WebSocketConnect(e.to_string()))?; + let (_, mut read) = ws.split(); + while let Some(msg) = read.next().await { + + if let Message::Binary(bin) = msg.map_err(|e| SubscriptionError::FailedToRecieveMessage(e.to_string()))?{ + yield rmp_serde::from_slice(&bin).map_err(|e| SubscriptionError::TiledEventDeserialize(e.to_string()))? + } + } + } + } + /// Create a new client for the given mock server #[cfg(test)] pub fn for_mock_server(server: &MockServer) -> Self { @@ -141,6 +184,27 @@ impl TiledClient { } } +#[derive(Debug, Clone)] +pub enum SubscriptionError { + WebSocketConnect(String), + FailedToRecieveMessage(String), + TiledEventDeserialize(String), +} +impl std::fmt::Display for SubscriptionError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + SubscriptionError::WebSocketConnect(err) => { + write!(f, "Failed to connect to tiled websocket: {}", err) + } + SubscriptionError::FailedToRecieveMessage(err) => { + write!(f, "Failed to recieve message from tiled websocket: {}", err) + } + SubscriptionError::TiledEventDeserialize(err) => { + write!(f, "Failed to convert tiled event into: {}", err) + } + } + } +} #[derive(Debug)] pub enum ClientError { InvalidPath(url::ParseError), diff --git a/src/handlers.rs b/src/handlers.rs index 705b1d6..d4347db 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -1,9 +1,9 @@ -use async_graphql::http::GraphiQLSource; -use async_graphql::{EmptyMutation, EmptySubscription, Schema}; -use async_graphql_axum::{GraphQLRequest, GraphQLResponse}; +use async_graphql::http::{ALL_WEBSOCKET_PROTOCOLS, GraphiQLSource}; +use async_graphql::{Data, EmptyMutation, Schema}; +use async_graphql_axum::{GraphQLProtocol, GraphQLRequest, GraphQLResponse, GraphQLWebSocket}; use axum::Extension; use axum::body::Body; -use axum::extract::{OptionalFromRequestParts, Path, State}; +use axum::extract::{OptionalFromRequestParts, Path, State, WebSocketUpgrade}; use axum::http::{HeaderMap, HeaderValue, StatusCode}; use axum::response::{Html, IntoResponse}; use reqwest::header::AUTHORIZATION; @@ -11,10 +11,11 @@ use tracing::info; use crate::clients::TiledClient; use crate::model::TiledQuery; +use crate::model::subscription::TiledSubscription; pub async fn graphql_handler( auth_token: Option, - schema: Extension>, + schema: Extension>, req: GraphQLRequest, ) -> GraphQLResponse { schema @@ -23,10 +24,41 @@ pub async fn graphql_handler( .into() } -pub async fn graphiql_handler(graphql_endpoint: Option) -> impl IntoResponse { +pub async fn graphql_ws_handler( + schema: Extension>, + protocol: GraphQLProtocol, + mut auth_token: Option, + websocket: WebSocketUpgrade, +) -> impl IntoResponse { + websocket + .protocols(ALL_WEBSOCKET_PROTOCOLS) + .on_upgrade(move |socket| { + GraphQLWebSocket::new(socket, schema.0, protocol) + .on_connection_init(|headers| async move { + let mut data = Data::default(); + + auth_token = headers + .get("Authorization") + .and_then(|v| v.as_str()) + .and_then(|auth| HeaderValue::from_str(auth).ok()) + .map(AuthHeader) + .or(auth_token); + + data.insert(auth_token); + Ok(data) + }) + .serve() + }) +} + +pub async fn graphiql_handler( + graphql_endpoint: Option, + subscription_endpoint: Option, +) -> impl IntoResponse { Html( GraphiQLSource::build() .endpoint(graphql_endpoint.as_deref().unwrap_or("/graphql")) + .subscription_endpoint(subscription_endpoint.as_deref().unwrap_or("/ws")) .finish(), ) } @@ -98,6 +130,7 @@ mod tests { fn app() -> Router { Router::new().route("/", get(auth_echo)) } + #[tokio::test] async fn auth_extract() { let app = app(); diff --git a/src/main.rs b/src/main.rs index 892cc27..a36a377 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use async_graphql::{EmptyMutation, EmptySubscription, Schema}; +use async_graphql::{EmptyMutation, Schema}; use axum::http::StatusCode; use axum::response::{Html, IntoResponse}; use axum::routing::{get, post}; @@ -26,8 +26,9 @@ use url::Url; use crate::clients::TiledClient; use crate::config::GlazedConfig; -use crate::handlers::{download_handler, graphiql_handler, graphql_handler}; +use crate::handlers::{download_handler, graphiql_handler, graphql_handler, graphql_ws_handler}; use crate::model::TiledQuery; +use crate::model::subscription::TiledSubscription; #[tokio::main] async fn main() -> Result<(), Box> { @@ -61,18 +62,27 @@ async fn serve(config: GlazedConfig) -> Result<(), Box> { .public_address .clone() .unwrap_or_else(|| Url::parse(&format!("http://{}", config.bind_address)).unwrap()); - let schema = Schema::build(TiledQuery, EmptyMutation, EmptySubscription) + let schema = Schema::build(TiledQuery, EmptyMutation, TiledSubscription) .data(RootAddress(public_address.clone())) .data(client.clone()) .finish(); let graphql_endpoint = config .public_address + .as_ref() .map(|u| u.join("graphql").unwrap().to_string()); + let subscription_endpoint = config + .public_address + .as_ref() + .map(|u| u.join("subscribe").unwrap().to_string()); let app = Router::new() .route("/graphql", post(graphql_handler).get(graphql_get_warning)) - .route("/graphiql", get(|| graphiql_handler(graphql_endpoint))) + .route("/subscribe", get(graphql_ws_handler)) + .route( + "/graphiql", + get(|| graphiql_handler(graphql_endpoint, subscription_endpoint)), + ) .route( "/status", get(Json(json!({"version": env!("CARGO_PKG_VERSION")}))), diff --git a/src/model.rs b/src/model.rs index b1cb3a9..bb9269e 100644 --- a/src/model.rs +++ b/src/model.rs @@ -4,6 +4,7 @@ pub(crate) mod container; pub(crate) mod event_stream; pub(crate) mod node; pub(crate) mod run; +pub(crate) mod subscription; pub(crate) mod table; use std::collections::HashMap; @@ -239,7 +240,7 @@ impl Run { #[cfg(test)] mod tests { - use async_graphql::{EmptyMutation, EmptySubscription, Schema, value}; + use async_graphql::{EmptyMutation, Schema, value}; use axum::http::HeaderValue; use httpmock::MockServer; use serde_json::json; @@ -247,9 +248,10 @@ mod tests { use crate::TiledQuery; use crate::clients::TiledClient; use crate::handlers::AuthHeader; + use crate::model::subscription::TiledSubscription; - fn build_schema(url: &str) -> Schema { - Schema::build(TiledQuery, EmptyMutation, EmptySubscription) + fn build_schema(url: &str) -> Schema { + Schema::build(TiledQuery, EmptyMutation, TiledSubscription) .data(Option::::None) .data(TiledClient::new(url.parse().unwrap())) .finish() @@ -320,7 +322,7 @@ mod tests { })); }) .await; - let schema = Schema::build(TiledQuery, EmptyMutation, EmptySubscription) + let schema = Schema::build(TiledQuery, EmptyMutation, TiledSubscription) .data(TiledClient::new(server.base_url().parse().unwrap())) .data(Some(AuthHeader::from(HeaderValue::from_static( "auth_value", diff --git a/src/model/subscription.rs b/src/model/subscription.rs new file mode 100644 index 0000000..c51783e --- /dev/null +++ b/src/model/subscription.rs @@ -0,0 +1,269 @@ +use async_graphql::{Context, SimpleObject, Subscription, Union}; +use futures_util::Stream; +use serde::Deserialize; +use serde_json::Value; + +use crate::clients::{SubscriptionError, TiledClient}; +use crate::handlers::AuthHeader; +use crate::model::array::{ArrayStructure, DataType}; + +pub struct TiledSubscription; + +#[Subscription] +impl TiledSubscription { + async fn events( + &self, + ctx: &Context<'_>, + ) -> impl Stream> { + let client = ctx.data::().unwrap(); + let headers = ctx + .data_opt::>() + .and_then(|a| a.as_ref()) + .map(|a| a.as_header_map()); + client.stream_events(None, headers) + } + + async fn node_events( + &self, + ctx: &Context<'_>, + node: String, + ) -> impl Stream> { + let client = ctx.data::().unwrap(); + let headers = ctx + .data_opt::>() + .and_then(|a| a.as_ref()) + .map(|a| a.as_header_map()); + client.stream_events(Some(node), headers) + } +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct ContainerSchema { + pub version: u32, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct ArraySchema { + pub version: u32, + pub data_type: DataType, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct TableSchema { + pub version: u32, + pub arrow_schema: String, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct ChildCreated { + pub sequence: u32, + pub timestamp: String, + pub key: String, + pub structure_family: String, + pub specs: Value, + pub metadata: Value, + pub data_sources: Vec, + pub access_blob: Value, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct ChildMetadataUpdated { + pub sequence: u32, + pub timestamp: String, + pub key: String, + pub specs: Value, + pub metadata: Value, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct StreamArrayData { + pub sequence: u32, + pub timestamp: String, + pub mimetype: String, + pub shape: Value, + pub offset: Option, + pub block: Option, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct ArrayRef { + pub sequence: u32, + pub timestamp: String, + pub data_source: Value, + pub patch: Option, + pub uri: Option, + pub shape: Value, + pub data_type: DataType, +} + +#[derive(Debug, Deserialize, SimpleObject)] +pub struct StreamTableData { + pub sequence: u32, + pub timestamp: String, + pub partition: Option, + pub append: bool, + pub arrow_schema: String, +} + +#[derive(Debug, PartialEq, Clone, Deserialize, SimpleObject)] +pub struct DataSource { + pub structure: ArrayStructure, +} + +#[derive(Union, Deserialize, Debug)] +#[serde(tag = "type")] +pub enum TiledEvent { + #[serde(rename = "container-schema")] + Container(ContainerSchema), + #[serde(rename = "array-schema")] + Array(ArraySchema), + #[serde(rename = "table-schema")] + Table(TableSchema), + #[serde(rename = "container-child-created")] + ChildCreated(ChildCreated), + #[serde(rename = "container-child-metadata-updated")] + ChildMetadataUpdated(ChildMetadataUpdated), + #[serde(rename = "array-data")] + ArrayData(StreamArrayData), + #[serde(rename = "array-ref")] + ArrayRef(ArrayRef), + #[serde(rename = "table-data")] + TableData(StreamTableData), +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use async_graphql::{EmptyMutation, Schema}; + use axum::Router; + use axum::extract::WebSocketUpgrade; + use axum::extract::ws::Message; + use axum::http::{HeaderMap, Uri}; + use axum::routing::any; + use futures_util::StreamExt; + use rstest::rstest; + use tokio::net::TcpListener; + use tokio::sync::mpsc; + use tokio::time; + use url::Url; + + use super::*; + use crate::clients::TiledClient; + use crate::model::TiledQuery; + + #[rstest] + #[case::with_auth(Some("Bearer test-token"))] + #[case::no_auth(None)] + #[tokio::test] + async fn test_stream_events_headers(#[case] auth_token: Option<&str>) { + let (tx, mut rx) = mpsc::channel(1); + + let app = Router::new().fallback(any( + move |headers: HeaderMap, ws: WebSocketUpgrade| async move { + let _ = tx.send(headers).await; + ws.on_upgrade(|_| async move {}) + }, + )); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let url = Url::parse(&format!("http://{addr}")).unwrap(); + let client = TiledClient::new(url); + + let mut headers = HeaderMap::new(); + if let Some(token) = auth_token { + headers.insert("Authorization", token.parse().unwrap()); + } + + let stream = client.stream_events(None, Some(headers)); + tokio::pin!(stream); + + let _ = time::timeout(Duration::from_secs(5), stream.next()).await; + + let captured_headers = time::timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("Headers should have been sent after stream is read above") + .expect("The header channel should remain open until headers are received"); + + if let Some(token) = auth_token { + assert_eq!(captured_headers.get("Authorization").unwrap(), token); + } else { + assert!(captured_headers.get("Authorization").is_none()); + } + } + + #[rstest::rstest] + #[case::root( + "subscription { events { ... on ContainerSchema { version } } }", + "/api/v1/stream/single/", + 42, + "events" + )] + #[case::node( + "subscription { nodeEvents(node: \"foo\") { ... on ContainerSchema { version } } }", + "/api/v1/stream/single/foo", + 43, + "nodeEvents" + )] + #[tokio::test] + async fn test_subscription_events( + #[case] query: &str, + #[case] expected_path: &str, + #[case] expected_version: u32, + #[case] result_key: &str, + ) { + let (tx, mut rx) = mpsc::channel(1); + + let app = Router::new().fallback(any(move |uri: Uri, ws: WebSocketUpgrade| async move { + let path = uri.path().to_string(); + let root = path.ends_with("/api/v1/stream/single/"); + ws.on_upgrade(move |mut socket| async move { + let version = if root { 42 } else { 43 }; + let event = serde_json::json!({ + "type": "container-schema", + "version": version + }); + let bin = rmp_serde::to_vec(&event).unwrap(); + socket.send(Message::Binary(bin.into())).await.unwrap(); + let _ = tx.send(path).await; + }) + })); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let url = Url::parse(&format!("http://{addr}")).unwrap(); + let client = TiledClient::new(url); + + let schema = Schema::build(TiledQuery, EmptyMutation, TiledSubscription) + .data(client) + .finish(); + + let stream = schema.execute_stream(query); + let mut stream = Box::pin(stream); + let res = time::timeout(Duration::from_secs(5), stream.next()) + .await + .expect("Timeout waiting for events") + .unwrap(); + + assert_eq!( + res.data, + async_graphql::value!({ result_key: { "version": expected_version } }) + ); + + let path = time::timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("Timeout waiting for path") + .unwrap(); + + assert_eq!(path, expected_path); + } +}