Add missing unit tests for subgraph batching

This commit adds a few more unit tests for verifying the functionality
of all helper methods and, where applicable, beefy methods pertinent to
the subgraph batching initiative.

apollo-router/src/batching.rs

@@ -222,3 +222,100 @@ impl Drop for Batch {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use hyper::body::to_bytes;
+    use tokio::sync::oneshot;
+
+    use super::Waiter;
+    use crate::graphql;
+    use crate::services::SubgraphRequest;
+    use crate::services::SubgraphResponse;
+    use crate::Context;
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn it_assembles_batch() {
+        let context = Context::new();
+
+        // Assemble a list of waiters for testing
+        let (receivers, waiters): (Vec<_>, Vec<_>) = (0..2)
+            .map(|index| {
+                let (tx, rx) = oneshot::channel();
+                let graphql_request = graphql::Request::fake_builder()
+                    .operation_name(format!("batch_test_{index}"))
+                    .query(format!("query batch_test {{ slot{index} }}"))
+                    .build();
+
+                (
+                    rx,
+                    Waiter::new(
+                        SubgraphRequest::fake_builder()
+                            .subgraph_request(
+                                http::Request::builder()
+                                    .body(graphql_request.clone())
+                                    .unwrap(),
+                            )
+                            .subgraph_name(format!("slot{index}"))
+                            .build(),
+                        graphql_request,
+                        context.clone(),
+                        tx,
+                    ),
+                )
+            })
+            .unzip();
+
+        // Try to assemble them
+        let (op_name, _context, request, txs) = Waiter::assemble_batch(waiters).await.unwrap();
+
+        // Make sure that the name of the entire batch is that of the first
+        assert_eq!(op_name, "batch_test_0");
+
+        // We should see the aggregation of all of the requests
+        let actual: Vec<graphql::Request> = serde_json::from_str(
+            &String::from_utf8(to_bytes(request.into_body()).await.unwrap().to_vec()).unwrap(),
+        )
+        .unwrap();
+
+        let expected: Vec<_> = (0..2)
+            .map(|index| {
+                graphql::Request::fake_builder()
+                    .operation_name(format!("batch_test_{index}"))
+                    .query(format!("query batch_test {{ slot{index} }}"))
+                    .build()
+            })
+            .collect();
+        assert_eq!(actual, expected);
+
+        // We should also have all of the correct senders and they should be linked to the correct waiter
+        // Note: We reverse the senders since they should be in reverse order when assembled
+        assert_eq!(txs.len(), receivers.len());
+        for (index, (tx, rx)) in Iterator::zip(txs.into_iter().rev(), receivers).enumerate() {
+            let data = serde_json_bytes::json!({
+                "data": {
+                    format!("slot{index}"): "valid"
+                }
+            });
+            let response = SubgraphResponse {
+                response: http::Response::builder()
+                    .body(graphql::Response::builder().data(data.clone()).build())
+                    .unwrap(),
+                context: Context::new(),
+            };
+
+            tx.send(Ok(response)).unwrap();
+
+            // We want to make sure that we don't hang the test if we don't get the correct message
+            let received = tokio::time::timeout(Duration::from_millis(10), rx)
+                .await
+                .unwrap()
+                .unwrap()
+                .unwrap();
+
+            assert_eq!(received.response.into_body().data, Some(data));
+        }
+    }
+}

apollo-router/src/services/subgraph_service.rs

@@ -2663,4 +2663,162 @@ mod tests {
 
         assert_eq!(resp.response.body(), &expected_resp);
     }
+
+    #[test]
+    fn it_gets_uri_details() {
+        let path = "https://example.com/path".parse().unwrap();
+        let (host, port, path) = super::get_uri_details(&path);
+
+        assert_eq!(host, "example.com");
+        assert_eq!(port, 443);
+        assert_eq!(path, "/path");
+    }
+
+    #[test]
+    fn it_converts_ok_http_to_graphql() {
+        let (parts, body) = http::Response::builder()
+            .status(StatusCode::OK)
+            .body(None)
+            .unwrap()
+            .into_parts();
+        let actual = super::http_response_to_graphql_response(
+            "test_service",
+            Ok(ContentType::ApplicationGraphqlResponseJson),
+            body,
+            &parts,
+        );
+
+        let expected = graphql::Response::builder().build();
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn it_converts_error_http_to_graphql() {
+        let (parts, body) = http::Response::builder()
+            .status(StatusCode::IM_A_TEAPOT)
+            .body(None)
+            .unwrap()
+            .into_parts();
+        let actual = super::http_response_to_graphql_response(
+            "test_service",
+            Ok(ContentType::ApplicationGraphqlResponseJson),
+            body,
+            &parts,
+        );
+
+        let expected = graphql::Response::builder()
+            .error(
+                super::FetchError::SubrequestHttpError {
+                    status_code: Some(418),
+                    service: "test_service".into(),
+                    reason: "418: I'm a teapot".into(),
+                }
+                .to_graphql_error(None),
+            )
+            .build();
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn it_converts_http_with_body_to_graphql() {
+        let mut json = serde_json::json!({
+            "data": {
+                "some_field": "some_value"
+            }
+        });
+
+        let (parts, body) = http::Response::builder()
+            .status(StatusCode::OK)
+            .body(Some(Ok(Bytes::from(json.to_string()))))
+            .unwrap()
+            .into_parts();
+
+        let actual = super::http_response_to_graphql_response(
+            "test_service",
+            Ok(ContentType::ApplicationGraphqlResponseJson),
+            body,
+            &parts,
+        );
+
+        let expected = graphql::Response::builder()
+            .data(json["data"].take())
+            .build();
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn it_converts_http_with_graphql_errors_to_graphql() {
+        let error = graphql::Error::builder()
+            .message("error was encountered for test")
+            .extension_code("SOME_EXTENSION")
+            .build();
+        let mut json = serde_json::json!({
+            "data": {
+                "some_field": "some_value",
+                "error_field": null,
+            },
+            "errors": [error],
+        });
+
+        let (parts, body) = http::Response::builder()
+            .status(StatusCode::OK)
+            .body(Some(Ok(Bytes::from(json.to_string()))))
+            .unwrap()
+            .into_parts();
+
+        let actual = super::http_response_to_graphql_response(
+            "test_service",
+            Ok(ContentType::ApplicationGraphqlResponseJson),
+            body,
+            &parts,
+        );
+
+        let expected = graphql::Response::builder()
+            .data(json["data"].take())
+            .error(error)
+            .build();
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn it_converts_error_http_with_graphql_errors_to_graphql() {
+        let error = graphql::Error::builder()
+            .message("error was encountered for test")
+            .extension_code("SOME_EXTENSION")
+            .build();
+        let mut json = serde_json::json!({
+            "data": {
+                "some_field": "some_value",
+                "error_field": null,
+            },
+            "errors": [error],
+        });
+
+        let (parts, body) = http::Response::builder()
+            .status(StatusCode::IM_A_TEAPOT)
+            .body(Some(Ok(Bytes::from(json.to_string()))))
+            .unwrap()
+            .into_parts();
+
+        let actual = super::http_response_to_graphql_response(
+            "test_service",
+            Ok(ContentType::ApplicationGraphqlResponseJson),
+            body,
+            &parts,
+        );
+
+        let expected = graphql::Response::builder()
+            .data(json["data"].take())
+            .error(
+                super::FetchError::SubrequestHttpError {
+                    status_code: Some(418),
+                    service: "test_service".into(),
+                    reason: "418: I'm a teapot".into(),
+                }
+                .to_graphql_error(None),
+            )
+            .error(error)
+            .build();
+        assert_eq!(actual, expected);
+    }
 }