From d8c700ce15fbf6a197207d613d9b08c37695ef91 Mon Sep 17 00:00:00 2001
From: Vasil Markoukin <susliko@users.noreply.github.com>
Date: Fri, 19 Apr 2024 18:02:42 +0300
Subject: [PATCH] Generate `anyOf` schema for `oneOfVariant`s with the same
 status code and content-type (#3703)

---
 .../openapi/EndpointToOperationResponse.scala | 22 +++++++-
 ...he_same_status_codes_and_content_types.yml | 53 +++++++++++++++++++
 .../docs/openapi/VerifyYamlOneOfTest.scala    | 19 ++++++-
 3 files changed, 92 insertions(+), 2 deletions(-)
 create mode 100644 docs/openapi-docs/src/test/resources/oneOf/expected_the_same_status_codes_and_content_types.yml

diff --git a/docs/openapi-docs/src/main/scala/sttp/tapir/docs/openapi/EndpointToOperationResponse.scala b/docs/openapi-docs/src/main/scala/sttp/tapir/docs/openapi/EndpointToOperationResponse.scala
index f4a8db6082..a08c2bc88b 100644
--- a/docs/openapi-docs/src/main/scala/sttp/tapir/docs/openapi/EndpointToOperationResponse.scala
+++ b/docs/openapi-docs/src/main/scala/sttp/tapir/docs/openapi/EndpointToOperationResponse.scala
@@ -84,7 +84,16 @@ private[openapi] class EndpointToOperationResponse(
 
     val description = bodies.headOption.flatMap { case (desc, _) => desc }.getOrElse(statusCodeDescriptions.headOption.getOrElse(""))
 
-    val content = bodies.flatMap { case (_, content) => content }.toListMap
+    val content = bodies
+      .flatMap { case (_, content) => content }
+      .foldLeft(ListMap.empty[String, Vector[MediaType]]) { case (acc, (ct, mt)) =>
+        acc.get(ct) match {
+          case Some(mts) => acc.updated(ct, mts :+ mt)
+          case None      => acc.updated(ct, Vector(mt))
+        }
+      }
+      .mapValues(mergeMediaTypesToAnyOf)
+      .toListMap
 
     if (bodies.nonEmpty || headers.nonEmpty) {
       Some(Response(description, headers.toListMap, content))
@@ -97,6 +106,17 @@ private[openapi] class EndpointToOperationResponse(
     }
   }
 
+  private def mergeMediaTypesToAnyOf(bodies: Vector[MediaType]): MediaType =
+    bodies.toSet.toList match {
+      case List(body) => body
+      case bodies =>
+        MediaType(
+          schema = Some(ASchema(anyOf = bodies.flatMap(_.schema))),
+          example = bodies.flatMap(_.example).headOption,
+          examples = bodies.flatMap(_.examples).toListMap
+        )
+    }
+
   private def collectBodies(outputs: List[EndpointOutput[_]]): List[(Option[String], ListMap[String, MediaType])] = {
     val forcedContentType = extractFixedContentType(outputs)
     outputs.flatMap(_.traverseOutputs {
diff --git a/docs/openapi-docs/src/test/resources/oneOf/expected_the_same_status_codes_and_content_types.yml b/docs/openapi-docs/src/test/resources/oneOf/expected_the_same_status_codes_and_content_types.yml
new file mode 100644
index 0000000000..289dac7efb
--- /dev/null
+++ b/docs/openapi-docs/src/test/resources/oneOf/expected_the_same_status_codes_and_content_types.yml
@@ -0,0 +1,53 @@
+openapi: 3.1.0
+info:
+  title: Fruits
+  version: '1.0'
+paths:
+  /:
+    get:
+      operationId: getRoot
+      responses:
+        '200':
+          description: not found
+          content:
+            application/json:
+              schema:
+                anyOf:
+                - $ref: '#/components/schemas/NotFound'
+                - $ref: '#/components/schemas/Unauthorized'
+        '204':
+          description: unknown
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/Unknown'
+components:
+  schemas:
+    NotFound:
+      title: NotFound
+      type: object
+      required:
+      - what
+      properties:
+        what:
+          type: string
+    Unauthorized:
+      title: Unauthorized
+      type: object
+      required:
+      - realm
+      properties:
+        realm:
+          type: string
+    Unknown:
+      title: Unknown
+      type: object
+      required:
+      - code
+      - msg
+      properties:
+        code:
+          type: integer
+          format: int32
+        msg:
+          type: string
diff --git a/docs/openapi-docs/src/test/scalajvm/sttp/tapir/docs/openapi/VerifyYamlOneOfTest.scala b/docs/openapi-docs/src/test/scalajvm/sttp/tapir/docs/openapi/VerifyYamlOneOfTest.scala
index c3feef4b5c..75e7a0dd49 100644
--- a/docs/openapi-docs/src/test/scalajvm/sttp/tapir/docs/openapi/VerifyYamlOneOfTest.scala
+++ b/docs/openapi-docs/src/test/scalajvm/sttp/tapir/docs/openapi/VerifyYamlOneOfTest.scala
@@ -59,7 +59,24 @@ class VerifyYamlOneOfTest extends AnyFunSuite with Matchers {
     actualYamlNoIndent shouldBe expectedYaml
   }
 
-  test("should support multiple the same status codes") {
+  test("should support multiple the same status codes and same content types") {
+    val expectedYaml = load("oneOf/expected_the_same_status_codes_and_content_types.yml")
+
+    val e = endpoint.out(
+      sttp.tapir.oneOf(
+        oneOfVariant(StatusCode.Ok, jsonBody[NotFound].description("not found")),
+        oneOfVariant(StatusCode.Ok, jsonBody[Unauthorized]),
+        oneOfVariant(StatusCode.NoContent, jsonBody[Unknown].description("unknown"))
+      )
+    )
+
+    val actualYaml = OpenAPIDocsInterpreter().toOpenAPI(e, Info("Fruits", "1.0")).toYaml
+    val actualYamlNoIndent = noIndentation(actualYaml)
+
+    actualYamlNoIndent shouldBe expectedYaml
+  }
+
+  test("should support multiple the same status codes and different content types") {
     val expectedYaml = load("oneOf/expected_the_same_status_codes.yml")
 
     implicit val unauthorizedTextPlainCodec: Codec[String, Unauthorized, CodecFormat.TextPlain] =