Add support for non-string discriminators (#1724)

RicoSuter · Nov 19, 2024 · 590632a · 590632a
1 parent e168e77
commit 590632a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 2 deletions.
diff --git a/src/NJsonSchema.Tests/Generation/SystemTextJson/SystemTextJsonInheritanceTests.cs b/src/NJsonSchema.Tests/Generation/SystemTextJson/SystemTextJsonInheritanceTests.cs
@@ -118,6 +118,46 @@ public async Task When_using_native_attributes_in_SystemTextJson_then_schema_is_
                     """.ReplaceLineEndings(), data);
         }
 
+        public class Dalmation : Dog
+        {
+            public string Foo { get; set; }
+        }
+
+        public class Poodle : Dog
+        {
+            public string Bar { get; set; }
+        }
+
+        [JsonDerivedType(typeof(Dalmation), 1)]
+        [JsonDerivedType(typeof(Poodle), 2)]
+        [JsonPolymorphic(TypeDiscriminatorPropertyName = "breed")]
+        public class Dog
+        {
+            public string Baz { get; set; }
+        }
+
+        [Fact]
+        public async Task When_using_native_attributes_and_integer_discriminator_in_SystemTextJson_then_schema_is_correct()
+        {
+            //// Act
+            var schema = JsonSchema.FromType<Dog>();
+            var data = schema.ToJson().ReplaceLineEndings();
+
+            //// Assert
+            Assert.NotNull(data);
+            Assert.Contains(@"""1"": """, data);
+            Assert.Contains(@"""2"": """, data);
+            Assert.Contains(
+                """
+                      "discriminator": {
+                        "propertyName": "breed",
+                        "mapping": {
+                          "1": "#/definitions/Dalmation",
+                          "2": "#/definitions/Poodle"
+                        }
+                      },
+                    """.ReplaceLineEndings(), data);
+        }
 #endif
     }
 }
diff --git a/src/NJsonSchema/Generation/JsonSchemaGenerator.cs b/src/NJsonSchema/Generation/JsonSchemaGenerator.cs
@@ -1141,9 +1141,9 @@ public SystemTextJsonInheritanceWrapper(string discriminatorName, dynamic[] json
 
             public string DiscriminatorName { get; }
 
-            public string GetDiscriminatorValue(Type type)
+            public string? GetDiscriminatorValue(Type type)
             {
-                return _jsonDerivedTypeAttributes.FirstOrDefault(a => a.DerivedType == type)?.TypeDiscriminator
+                return _jsonDerivedTypeAttributes.FirstOrDefault(a => a.DerivedType == type)?.TypeDiscriminator?.ToString()
                     ?? throw new InvalidOperationException($"Discriminator value for {type.FullName} not found.");
             }
         }