From 1e6ae23cd891ca91beb28a9cec71812bd7b8b24b Mon Sep 17 00:00:00 2001
From: Sunli <scott_s829@163.com>
Date: Sat, 6 Mar 2021 08:21:44 +0800
Subject: [PATCH] Add test for Federation entity lookup with DataLoader.

---
 tests/federation.rs | 105 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 105 insertions(+)

diff --git a/tests/federation.rs b/tests/federation.rs
index 708cbabf..2b6308ab 100644
--- a/tests/federation.rs
+++ b/tests/federation.rs
@@ -1,5 +1,9 @@
 #![allow(unreachable_code)]
 
+use std::collections::HashMap;
+use std::convert::Infallible;
+
+use async_graphql::dataloader::{DataLoader, Loader};
 use async_graphql::*;
 
 #[async_std::test]
@@ -146,3 +150,104 @@ pub async fn test_federation() {
         })
     );
 }
+
+#[async_std::test]
+pub async fn test_find_entity_with_context() {
+    struct MyLoader;
+
+    #[async_trait::async_trait]
+    impl Loader<ID> for MyLoader {
+        type Value = MyObj;
+        type Error = Infallible;
+
+        async fn load(&self, keys: &[ID]) -> Result<HashMap<ID, Self::Value>, Self::Error> {
+            Ok(keys
+                .iter()
+                .filter(|id| id.as_str() != "999")
+                .map(|id| {
+                    (
+                        id.clone(),
+                        MyObj {
+                            id: id.clone(),
+                            value: 999,
+                        },
+                    )
+                })
+                .collect())
+        }
+    }
+
+    #[derive(Clone, SimpleObject)]
+    struct MyObj {
+        id: ID,
+        value: i32,
+    }
+
+    struct QueryRoot;
+
+    #[Object]
+    impl QueryRoot {
+        #[graphql(entity)]
+        async fn find_user_by_id(&self, ctx: &Context<'_>, id: ID) -> FieldResult<MyObj> {
+            let loader = ctx.data_unchecked::<DataLoader<MyLoader>>();
+            loader
+                .load_one(id)
+                .await
+                .unwrap()
+                .ok_or_else(|| "Not found".into())
+        }
+    }
+
+    let schema = Schema::build(QueryRoot, EmptyMutation, EmptySubscription)
+        .data(DataLoader::new(MyLoader))
+        .finish();
+    let query = r#"{
+            _entities(representations: [
+                {__typename: "MyObj", id: "1"},
+                {__typename: "MyObj", id: "2"},
+                {__typename: "MyObj", id: "3"},
+                {__typename: "MyObj", id: "4"}
+            ]) {
+                __typename
+                ... on MyObj {
+                    id
+                    value
+                }
+            }
+        }"#;
+    assert_eq!(
+        schema.execute(query).await.into_result().unwrap().data,
+        value!({
+            "_entities": [
+                {"__typename": "MyObj", "id": "1", "value": 999 },
+                {"__typename": "MyObj", "id": "2", "value": 999 },
+                {"__typename": "MyObj", "id": "3", "value": 999 },
+                {"__typename": "MyObj", "id": "4", "value": 999 },
+            ]
+        })
+    );
+
+    let query = r#"{
+            _entities(representations: [
+                {__typename: "MyObj", id: "999"}
+            ]) {
+                __typename
+                ... on MyObj {
+                    id
+                    value
+                }
+            }
+        }"#;
+    assert_eq!(
+        schema.execute(query).await.into_result().unwrap_err(),
+        vec![ServerError {
+            message: "Not found".to_string(),
+            locations: vec![Pos {
+                line: 2,
+                column: 13
+            }],
+            path: vec![PathSegment::Field("_entities".to_owned())],
+            extensions: None,
+        }]
+    );
+}