From e43e92cb9e1a06bfe78f37ac7e7c21216ea22170 Mon Sep 17 00:00:00 2001
From: Vitor Hideyoshi <vitor.h.n.batista@gmail.com>
Date: Sun, 23 Nov 2025 20:03:19 -0300
Subject: [PATCH] feat: minor adjustments to oneOf and adds tests for examples
 in allOf, oneOf, anyOf

---
 jambo/parser/oneof_type_parser.py      |  3 +-
 tests/parser/test_anyof_type_parser.py | 28 ++++++++++-
 tests/parser/test_oneof_type_parser.py | 68 ++++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/jambo/parser/oneof_type_parser.py b/jambo/parser/oneof_type_parser.py
index 4713d9c..11092ea 100644
--- a/jambo/parser/oneof_type_parser.py
+++ b/jambo/parser/oneof_type_parser.py
@@ -45,8 +45,7 @@ class OneOfTypeParser(GenericTypeParser):
         # they were added by OpenAPI and not all implementations may support them,
         # and they do not always generate a model one-to-one to the Pydantic model
         # TL;DR: Discriminators were added by OpenAPI and not a Official JSON Schema feature
-        discriminator = properties.get("discriminator")
-        if discriminator is not None:
+        if (discriminator := properties.get("discriminator")) is not None:
             validated_type = self._build_type_one_of_with_discriminator(
                 subfield_types, discriminator
             )
diff --git a/tests/parser/test_anyof_type_parser.py b/tests/parser/test_anyof_type_parser.py
index 811bc1b..7f26a80 100644
--- a/tests/parser/test_anyof_type_parser.py
+++ b/tests/parser/test_anyof_type_parser.py
@@ -99,7 +99,33 @@ class TestAnyOfTypeParser(TestCase):
         with self.assertRaises(InvalidSchemaException):
             AnyOfTypeParser().from_properties("placeholder", properties)
 
-    def test_any_of_with_examples(self):
+    def test_anyof_with_examples(self):
+        """
+        Tests the AnyOfTypeParser with a string or int type and examples.
+        """
+
+        properties = {
+            "anyOf": [
+                {
+                    "type": "string",
+                    "examples": ["example string"],
+                },
+                {
+                    "type": "integer",
+                    "examples": [123],
+                },
+            ],
+        }
+
+        parsed_type, _ = AnyOfTypeParser().from_properties("placeholder", properties)
+
+        type_1, type_2 = get_args(parsed_type)
+
+        self.assertEqual(get_args(type_1)[1].examples, ["example string"])
+
+        self.assertEqual(get_args(type_2)[1].examples, [123])
+
+    def test_any_of_with_root_examples(self):
         """
         Tests the AnyOfTypeParser with a string or int type and examples.
         """
diff --git a/tests/parser/test_oneof_type_parser.py b/tests/parser/test_oneof_type_parser.py
index 14bf942..8fa3900 100644
--- a/tests/parser/test_oneof_type_parser.py
+++ b/tests/parser/test_oneof_type_parser.py
@@ -532,3 +532,71 @@ class TestOneOfTypeParser(TestCase):
         # Invalid: Wrong properties for the type
         with self.assertRaises(ValidationError):
             Model(shape={"type": "circle", "width": 10})
+
+    def test_oneof_with_examples(self):
+        schema = {
+            "title": "ExampleTest",
+            "type": "object",
+            "properties": {
+                "value": {
+                    "oneOf": [
+                        {
+                            "type": "string",
+                            "examples": ["example1", "example2"],
+                        },
+                        {
+                            "type": "integer",
+                            "examples": [1, 2, 3],
+                        },
+                    ]
+                }
+            },
+            "required": ["value"],
+        }
+
+        Model = SchemaConverter.build(schema)
+
+        # Since Pydantic does not natively support oneOf and the validation
+        # is done via a custom validator, the `value` is represented using `anyOf`
+        model_schema = Model.model_json_schema()
+
+        self.assertEqual(
+            model_schema["properties"]["value"]["anyOf"][0]["examples"],
+            ["example1", "example2"],
+        )
+
+        self.assertEqual(
+            model_schema["properties"]["value"]["anyOf"][1]["examples"],
+            [1, 2, 3],
+        )
+
+    def test_oneof_with_root_examples(self):
+        schema = {
+            "title": "ExampleTest",
+            "type": "object",
+            "properties": {
+                "value": {
+                    "oneOf": [
+                        {
+                            "type": "string",
+                        },
+                        {
+                            "type": "integer",
+                        },
+                    ],
+                    "examples": ["example1", 2],
+                }
+            },
+            "required": ["value"],
+        }
+
+        Model = SchemaConverter.build(schema)
+
+        # Since Pydantic does not natively support oneOf and the validation
+        # is done via a custom validator, the `value` is represented using `anyOf`
+        model_schema = Model.model_json_schema()
+
+        self.assertEqual(
+            model_schema["properties"]["value"]["examples"],
+            ["example1", 2],
+        )