forked from vectordotdev/vector
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest.rs
More file actions
58 lines (52 loc) · 1.63 KB
/
Copy pathtest.rs
File metadata and controls
58 lines (52 loc) · 1.63 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
use crate::{
config::{SinkConfig, SinkContext},
Error,
};
use bytes::Bytes;
use futures::{FutureExt, TryFutureExt};
use hyper::{
service::{make_service_fn, service_fn},
Body, Request, Response, Server,
};
use serde::Deserialize;
use stream_cancel::{Trigger, Tripwire};
use tokio::sync::mpsc;
pub fn load_sink<T>(config: &str) -> crate::Result<(T, SinkContext)>
where
for<'a> T: Deserialize<'a> + SinkConfig,
{
let sink_config: T = toml::from_str(config)?;
let cx = SinkContext::new_test();
Ok((sink_config, cx))
}
pub fn build_test_server(
addr: std::net::SocketAddr,
) -> (
mpsc::Receiver<(http::request::Parts, Bytes)>,
Trigger,
impl std::future::Future<Output = Result<(), ()>>,
) {
let (tx, rx) = mpsc::channel(100);
let service = make_service_fn(move |_| {
let tx = tx.clone();
async {
Ok::<_, Error>(service_fn(move |req: Request<Body>| {
let mut tx = tx.clone();
async {
let (parts, body) = req.into_parts();
tokio::spawn(async move {
let bytes = hyper::body::to_bytes(body).await.unwrap();
tx.send((parts, bytes)).await.unwrap();
});
Ok::<_, Error>(Response::new(Body::empty()))
}
}))
}
});
let (trigger, tripwire) = Tripwire::new();
let server = Server::bind(&addr)
.serve(service)
.with_graceful_shutdown(tripwire.then(crate::stream::tripwire_handler))
.map_err(|error| panic!("Server error: {}", error));
(rx, trigger, server)
}