From f07251811167a3d84a704ddde8d980c11c1057ad Mon Sep 17 00:00:00 2001
From: Nick Iodice <97116136+niodice@users.noreply.github.com>
Date: Thu, 12 Dec 2024 16:03:06 -0700
Subject: [PATCH] Adhere to specification for schema extension (#2433)

---
 .../caliban/rendering/DocumentRenderer.scala  |  3 +-
 .../test/scala/caliban/RenderingSpec.scala    | 30 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/core/src/main/scala/caliban/rendering/DocumentRenderer.scala b/core/src/main/scala/caliban/rendering/DocumentRenderer.scala
index 78c76a70bc..3a83940325 100644
--- a/core/src/main/scala/caliban/rendering/DocumentRenderer.scala
+++ b/core/src/main/scala/caliban/rendering/DocumentRenderer.scala
@@ -327,7 +327,8 @@ object DocumentRenderer extends Renderer[Document] {
       definition match {
         case SchemaDefinition(directives, query, mutation, subscription, description) =>
           val hasTypes    = query.nonEmpty || mutation.nonEmpty || subscription.nonEmpty
-          val isExtension = directives.nonEmpty && !hasTypes
+          val isExtension =
+            (!hasTypes && directives.nonEmpty) || (query.isEmpty && (mutation.nonEmpty || subscription.nonEmpty))
           var first       = true
 
           def renderOp(name: String, op: Option[String]): Unit =
diff --git a/core/src/test/scala/caliban/RenderingSpec.scala b/core/src/test/scala/caliban/RenderingSpec.scala
index 2281fc1165..889282a27b 100644
--- a/core/src/test/scala/caliban/RenderingSpec.scala
+++ b/core/src/test/scala/caliban/RenderingSpec.scala
@@ -17,6 +17,7 @@ import caliban.schema.Annotations.GQLOneOfInput
 import caliban.schema.Schema.auto._
 import caliban.schema.ArgBuilder.auto._
 import caliban.schema.{ ArgBuilder, Schema }
+import zio.stream.ZStream
 import zio.{ IO, ZIO }
 import zio.test.Assertion._
 import zio.test._
@@ -101,7 +102,7 @@ object RenderingSpec extends ZIOSpecDefault {
         assertTrue(graphQL(InvalidSchemas.resolverEmpty).render.trim == "")
       },
       test(
-        "it should render a schema extension with schema directives even if no queries, mutations, or subscription"
+        "it should render a schema extension with directives only"
       ) {
         val renderedType =
           graphQL(InvalidSchemas.resolverEmpty, schemaDirectives = List(SchemaDirectives.Link)).render.trim
@@ -109,6 +110,33 @@ object RenderingSpec extends ZIOSpecDefault {
           renderedType == """extend schema @link(url: "https://example.com", import: ["@key", {name: "@provides", as: "@self"}])"""
         )
       },
+      test("it should render a schema extension with directives and a mutation") {
+        val resolver     = RootResolver(
+          Option.empty[Unit],
+          Some(MutationIO(_ => ZIO.unit)),
+          Option.empty[Unit]
+        )
+        val renderedType = graphQL(resolver, schemaDirectives = List(SchemaDirectives.Link)).render.trim
+        assertTrue(renderedType.startsWith("extend schema"))
+      },
+      test("it should render a schema extension with directives and a subscription") {
+        val resolver     = RootResolver(
+          Option.empty[Unit],
+          Option.empty[Unit],
+          Some(SubscriptionIO(ZStream.empty))
+        )
+        val renderedType = graphQL(resolver, schemaDirectives = List(SchemaDirectives.Link)).render.trim
+        assertTrue(renderedType.startsWith("extend schema"))
+      },
+      test("it should render a schema extension with a subscription and mutation but no directives") {
+        val resolver     = RootResolver(
+          Option.empty[Unit],
+          Some(MutationIO(_ => ZIO.unit)),
+          Some(SubscriptionIO(ZStream.empty))
+        )
+        val renderedType = graphQL(resolver).render.trim
+        assertTrue(renderedType.startsWith("extend schema"))
+      },
       test("it should render object arguments in type directives") {
         val testType     = __Type(
           __TypeKind.OBJECT,