test lambda function end-to-end

davidbarsky · davidbarsky · commit 1b95ac1f5e1c · 2019-12-25T16:51:40.000-05:00
diff --git a/lambda/src/client.rs b/lambda/src/client.rs
@@ -10,6 +10,7 @@ use http::{
     HeaderValue, Method, Request, Response, StatusCode, Uri,
 };
 use hyper::Body;
+use serde_json::json;
 use std::{
     convert::{TryFrom, TryInto},
     future::Future,
@@ -70,9 +71,9 @@ where
     }
 }
 
-pub struct NextEventSvc;
+pub struct EndpointSvc;
 
-impl Service<Request<Body>> for NextEventSvc {
+impl Service<Request<Body>> for EndpointSvc {
     type Response = Response<Body>;
     type Error = crate::Err;
     type Future = Fut<'static, Result<Self::Response, Self::Error>>;
@@ -101,13 +102,14 @@ async fn next_event(req: Request<Body>) -> Result<Response<Body>, Err> {
     let path = "/2018-06-01/runtime/invocation/next";
     assert_eq!(req.method(), Method::GET);
     assert_eq!(req.uri().path_and_query().unwrap(), &PathAndQuery::from_static(path));
+    let body = json!({"message": "hello"});
 
     let rsp = NextEventResponse {
         request_id: "8476a536-e9f4-11e8-9739-2dfe598c3fcd",
         deadline: 1542409706888,
         arn: "arn:aws:lambda:us-east-2:123456789012:function:custom-runtime",
         trace_id: "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419",
-        body: vec![],
+        body: serde_json::to_vec(&body)?,
     };
     rsp.into_rsp()
 }
@@ -144,7 +146,7 @@ async fn event_err(req: Request<Body>) -> Result<Response<Body>, Err> {
 pub struct MakeSvc;
 
 impl<T> Service<T> for MakeSvc {
-    type Response = NextEventSvc;
+    type Response = EndpointSvc;
     type Error = std::io::Error;
     type Future = future::Ready<Result<Self::Response, Self::Error>>;
 
@@ -153,15 +155,17 @@ impl<T> Service<T> for MakeSvc {
     }
 
     fn call(&mut self, _: T) -> Self::Future {
-        future::ok(NextEventSvc)
+        future::ok(EndpointSvc)
     }
 }
 
 #[cfg(test)]
 mod endpoint_tests {
     use super::{Client, MakeSvc};
     use crate::{
+        context, handler_fn,
         requests::{EventCompletionRequest, EventErrorRequest, IntoRequest, NextEventRequest},
+        run_simulated,
         types::Diagnostic,
         Err,
     };
@@ -254,4 +258,30 @@ mod endpoint_tests {
         future::try_select(server, client).await.unwrap();
         Ok(())
     }
+
+    #[tokio::test]
+    async fn run_end_to_end() -> Result<(), Err> {
+        use serde_json::Value;
+        let (listener, addr) = setup()?;
+        let url = format!("http://{}/", addr);
+
+        let server = tokio::spawn(async move {
+            let svc = hyper::Server::from_tcp(listener)?.serve(MakeSvc);
+            svc.await
+        });
+
+        async fn handler(s: Value) -> Result<Value, Err> {
+            let ctx = context();
+            assert!(ctx.xray_trace_id.is_some());
+            Ok(s)
+        }
+        let handler = handler_fn(handler);
+
+        let handler = tokio::spawn(async move {
+            run_simulated(handler, &url).await?;
+            Ok::<(), Err>(())
+        });
+        future::try_select(server, handler).await.unwrap();
+        Ok(())
+    }
 }
diff --git a/lambda/src/lib.rs b/lambda/src/lib.rs
@@ -33,7 +33,6 @@
 //! }
 //! ```
 pub use crate::types::LambdaCtx;
-use bytes::buf::BufExt;
 use client::Client;
 use http::{Request, Response};
 use hyper::Body;
@@ -129,10 +128,6 @@ pub trait Handler<A, B> {
     /// The future response value of this handler.
     type Fut: Future<Output = Result<B, Self::Err>>;
     /// Process the incoming event and return the response asynchronously.
-    ///
-    /// # Arguments
-    /// * `event` - The data received in the invocation request
-    /// * `ctx` - The context for the current invocation
     fn call(&mut self, event: A) -> Self::Fut;
 }
 
@@ -193,7 +188,23 @@ where
     let config = Config::from_env()?;
     let client = Client::with(&config.endpoint, hyper::Client::new())?;
     let mut exec = Executor { client };
-    exec.run(&mut handler).await?;
+    exec.run(&mut handler, false).await?;
+
+    Ok(())
+}
+
+/// Runs the lambda function almost entirely in-memory. This is mean for easier testing.
+pub async fn run_simulated<A, B, F>(handler: F, url: &str) -> Result<(), Err>
+where
+    F: Handler<A, B>,
+    <F as Handler<A, B>>::Err: fmt::Debug,
+    A: for<'de> Deserialize<'de>,
+    B: Serialize,
+{
+    let mut handler = handler;
+    let client = Client::with(url, hyper::Client::new())?;
+    let mut exec = Executor { client };
+    exec.run(&mut handler, true).await?;
 
     Ok(())
 }
@@ -207,24 +218,24 @@ where
     S: Service<Request<Body>, Response = Response<Body>>,
     <S as Service<Request<Body>>>::Error: Into<Err> + Send + Sync + 'static + std::error::Error,
 {
-    async fn run<A, B, F>(&mut self, handler: &mut F) -> Result<(), Err>
+    async fn run<A, B, F>(&mut self, handler: &mut F, once: bool) -> Result<(), Err>
     where
         F: Handler<A, B>,
         <F as Handler<A, B>>::Err: fmt::Debug,
         A: for<'de> Deserialize<'de>,
         B: Serialize,
     {
         let client = &mut self.client;
+        // todo: refactor this into a stream so that the `once` boolean can be replaced
+        // a `.take(n).await` combinator if we want to run this once, if `n` is `1`.
         loop {
             let req = NextEventRequest.into_req()?;
             let event = client.call(req).await?;
             let (parts, body) = event.into_parts();
 
-            let mut ctx = LambdaCtx::try_from(parts.headers)?;
-            ctx.env_config = Config::from_env()?;
-
-            let body = hyper::body::aggregate(body).await?;
-            let body = serde_json::from_reader(body.reader())?;
+            let ctx = LambdaCtx::try_from(parts.headers)?;
+            let body = hyper::body::to_bytes(body).await?;
+            let body = serde_json::from_slice(&body)?;
 
             let request_id = &ctx.request_id.clone();
             let f = WithTaskLocal {
@@ -244,6 +255,9 @@ where
                 .into_req()?,
             };
             client.call(req).await?;
+            if once {
+                break Ok(());
+            }
         }
     }
 }
diff --git a/lambda/src/types.rs b/lambda/src/types.rs
@@ -130,8 +130,14 @@ impl TryFrom<HeaderMap> for LambdaCtx {
                 .expect("Missing deadline"),
             invoked_function_arn: headers["lambda-runtime-invoked-function-arn"]
                 .to_str()
-                .expect("Missing arn")
+                .expect("Missing arn; this is a bug")
                 .to_owned(),
+            xray_trace_id: headers.get("lambda-runtime-trace-id").map(|header| {
+                header
+                    .to_str()
+                    .expect("Invalid XRayTraceID sent by Lambda; this is a bug")
+                    .to_owned()
+            }),
             ..Default::default()
         };
         Ok(ctx)
