Resolve schema modules in discriminator mapping

lib/open_api_spex/schema_resolver.ex

@@ -2,6 +2,8 @@ defmodule OpenApiSpex.SchemaResolver do
   @moduledoc """
   Internal module used to resolve `OpenApiSpex.Schema` structs from atoms.
   """
+  alias OpenApiSpex.Discriminator
+
   alias OpenApiSpex.{
     Components,
     MediaType,
@@ -228,6 +230,9 @@ defmodule OpenApiSpex.SchemaResolver do
     {properties, schemas} =
       resolve_schema_modules_from_schema_properties(schema.properties, schemas)
 
+    {discriminator, schemas} =
+      resolve_schema_modules_from_discriminator(schema.discriminator, schemas)
+
     schema = %{
       schema
       | allOf: all_of,
@@ -236,7 +241,8 @@ defmodule OpenApiSpex.SchemaResolver do
         not: not_schema,
         items: items,
         additionalProperties: additional,
-        properties: properties
+        properties: properties,
+        discriminator: discriminator
     }
 
     {schema, schemas}
@@ -257,4 +263,23 @@ defmodule OpenApiSpex.SchemaResolver do
   defp resolve_schema_modules_from_schema_properties(properties, _schemas) do
     raise "Expected :properties to be a map. Got: #{inspect(properties)}"
   end
+
+  defp resolve_schema_modules_from_discriminator(
+         discriminator = %Discriminator{mapping: mapping = %{}},
+         schemas
+       ) do
+    {mapping, schemas} =
+      Enum.map_reduce(mapping, schemas, fn
+        {key, module}, schemas when is_atom(module) ->
+          {%Reference{"$ref": path}, schemas} = resolve_schema_modules_from_schema(module, schemas)
+          {{key, path}, schemas}
+
+        {key, path}, schemas ->
+          {{key, path}, schemas}
+      end)
+
+    {%{discriminator | mapping: Map.new(mapping)}, schemas}
+  end
+
+  defp resolve_schema_modules_from_discriminator(disciminator, schemas), do: {disciminator, schemas}
 end

test/schema_resolver_test.exs

@@ -134,6 +134,25 @@ defmodule OpenApiSpex.SchemaResolverTest do
               }
             }
           }
+        },
+        "/api/appointsments" => %PathItem{
+          post: %Operation{
+            description: "Create a new pet appointment",
+            operationId: "PetAppointmentController.create",
+            requestBody: %RequestBody{
+              required: true,
+              content: %{
+                "application/json" => %MediaType{
+                  schema: OpenApiSpexTest.Schemas.PetAppointmentRequest
+                }
+              }
+            },
+            responses: %{
+              201 => %Response{
+                description: "Appointment created"
+              }
+            }
+          }
         }
       }
     }
@@ -149,14 +168,23 @@ defmodule OpenApiSpex.SchemaResolverTest do
     assert %Reference{"$ref": "#/components/schemas/UserRequest"} =
              resolved.paths["/api/users"].post.requestBody.content["application/json"].schema
 
+    assert "#/components/schemas/TrainingAppointment" =
+             resolved.components.schemas["PetAppointmentRequest"].discriminator.mapping["training"]
+
+    assert "#/components/schemas/GroomingAppointment" =
+             resolved.components.schemas["PetAppointmentRequest"].discriminator.mapping["grooming"]
+
     assert %{
              "UserRequest" => %Schema{},
              "UserResponse" => %Schema{},
              "User" => %Schema{},
              "UserSubscribeRequest" => %Schema{},
              "PaymentDetails" => %Schema{},
              "CreditCardPaymentDetails" => %Schema{},
-             "DirectDebitPaymentDetails" => %Schema{}
+             "DirectDebitPaymentDetails" => %Schema{},
+             "PetAppointmentRequest" => %Schema{},
+             "TrainingAppointment" => %Schema{},
+             "GroomingAppointment" => %Schema{}
            } = resolved.components.schemas
 
     get_friends = resolved.paths["/api/users/{id}/friends"].get

test/support/schemas.ex

@@ -645,8 +645,8 @@ defmodule OpenApiSpexTest.Schemas do
       discriminator: %OpenApiSpex.Discriminator{
         propertyName: "appointment_type",
         mapping: %{
-          "training" => "TrainingAppointment",
-          "grooming" => "GroomingAppointment"
+          "training" => TrainingAppointment,
+          "grooming" => GroomingAppointment
         }
       }
     })