Modify the fold function to match roots

Author: hmy3743

src/com/walmartlabs/lacinia/federation.clj

@@ -314,21 +314,24 @@
 
 (defn ^:private fold-queries
   [{:keys [queries] :as schema}]
-  (cond
-    (map? queries) (update-in schema [:objects :Query :fields] merge queries)
-    :else schema))
+  (let [query (get-in schema [:roots :query] :Query)]
+    (cond
+     (map? queries) (update-in schema [:objects query :fields] merge queries)
+     :else schema)))
 
 (defn ^:private fold-mutations
   [{:keys [mutations] :as schema}]
-  (cond
-    (map? mutations) (update-in schema [:objects :Mutation :fields] merge mutations)
-    :else schema))
+  (let [mutation (get-in schema [:roots :mutation] :Mutation)]
+    (cond
+     (map? mutations) (update-in schema [:objects mutation :fields] merge mutations)
+     :else schema)))
 
 (defn ^:private fold-subscriptions
   [{:keys [subscriptions] :as schema}]
-  (cond
-    (map? subscriptions) (update-in schema [:objects :Subscription :fields] merge subscriptions)
-    :else schema))
+  (let [subscription (get-in schema [:roots :subscription] :Subscription)]
+    (cond
+     (map? subscriptions) (update-in schema [:objects subscription :fields] merge subscriptions)
+     :else schema)))
 
 (defn generate-sdl
   "Translate the edn lacinia schema to the SDL schema."
@@ -355,7 +358,8 @@
                 :directive-defs (edn-directive-defs->sdl-directives val)
                 :roots (edn-roots->sdl-schema val)
                 "")))
-       (join "\n\n")))
+       (join "\n\n")
+       clojure.string/trim))
 
 (defn inject-federation
   "Called after SDL parsing to extend the input schema

test/com/walmartlabs/lacinia/federation_tests.clj

@@ -269,15 +269,15 @@ query($reps : [_Any!]!) {
       (is (= #{"Stuff"} union-names)))))
 
 (deftest edn-schema->sdl-schema
-  (let [sample-schema '{:roots {:query :query
-                                :mutation :mutation}
+  (let [sample-edn-1 '{:roots {:query :MyQuery
+                                :mutation :Mutation}
                         :interfaces
                         {:Node
                          {:fields
                           {:id
                            {:type (non-null ID)}}}}
                         :objects
-                        {:Query
+                        {:MyQuery
                          {:fields
                           {:todo
                            {:type :Todo
@@ -332,8 +332,18 @@ query($reps : [_Any!]!) {
                            :resolvable
                            {:type Boolean :default-value true}}}
                          :external
-                         {:locations #{:field-definition}}}}]
-    (is (= (-> sample-schema generate-sdl parse-schema) sample-schema))))
+                         {:locations #{:field-definition}}}}
+        sample-edn-2 '{:queries
+                       {:node
+                        {:description "node query"
+                         :type Node
+                         :args {:id {:type (non-null ID)}}}}
+                       :roots
+                       {:query :CustomQuery}}
+        sample-sdl-2 "schema {\n  query: CustomQuery\n}\n\ntype CustomQuery{\n  \"\"\"\n  node query\n  \"\"\"\n  node(id: ID!): Node\n}"]
+    
+    (is (= (-> sample-edn-1 generate-sdl parse-schema) sample-edn-1))
+    (is (= (generate-sdl sample-edn-2) sample-sdl-2))))
 
 (deftest only-edn-schama-essential
   (let [edn (-> "dev-resources/edn-federation.edn" slurp read-string)