test: add coverage for composition variant propagation and rewriting

Author: damaz91

tests/test_codegen_pipeline.py

@@ -347,6 +347,49 @@ def test_array_variant_preserves_root_and_filters_nested_objects(
         self.assertEqual(set(item_schema["required"]), {"amount", "label"})
         self.assertEqual(variant["title"], "Totals Create Request")
 
+    def test_composition_variant_rewrites_refs(self) -> None:
+        """Composition variants (oneOf/anyOf/allOf) rewrite refs to variants."""
+        schema = {
+            "$id": "https://ucp.dev/schemas/poly.json",
+            "title": "Poly",
+            "oneOf": [{"$ref": "child_a.json"}, {"$ref": "child_b.json"}],
+            "allOf": [{"$ref": "parent.json"}],
+            "anyOf": [{"$ref": "other.json"}],
+        }
+        file_path = Path("/schemas/poly.json")
+        child_a_path = str((file_path.parent / "child_a.json").resolve())
+        child_b_path = str((file_path.parent / "child_b.json").resolve())
+        parent_path = str((file_path.parent / "parent.json").resolve())
+        other_path = str((file_path.parent / "other.json").resolve())
+
+        variant_needs = {
+            child_a_path: {"create"},
+            child_b_path: {"create"},
+            parent_path: {"create"},
+            other_path: {"create"},
+        }
+
+        variant = preprocess_schemas._create_single_variant(
+            schema,
+            "create",
+            "poly",
+            file_path,
+            variant_needs,
+        )
+
+        self.assertEqual(
+            variant["oneOf"][0]["$ref"], "child_a_create_request.json"
+        )
+        self.assertEqual(
+            variant["oneOf"][1]["$ref"], "child_b_create_request.json"
+        )
+        self.assertEqual(
+            variant["allOf"][0]["$ref"], "parent_create_request.json"
+        )
+        self.assertEqual(
+            variant["anyOf"][0]["$ref"], "other_create_request.json"
+        )
+
     def test_generate_variants_writes_operation_specific_files(self) -> None:
         """Variant generation writes one filtered file per operation."""
         schema = {
@@ -473,6 +516,26 @@ def test_variant_needs_propagate_transitively_and_respect_omit(
         self.assertEqual(variant_needs[child_path], {"create"})
         self.assertEqual(variant_needs[grandchild_path], {"create"})
 
+    def test_variant_needs_propagate_through_composition_keywords(self) -> None:
+        """Variant dependencies propagate unconditionally through oneOf/anyOf/allOf/items."""
+        parent_path = "/schemas/parent.json"
+        child_path = "/schemas/child.json"
+        schemas = {
+            parent_path: {"oneOf": [{"$ref": "child.json"}]},
+            child_path: {"properties": {}},
+        }
+        schema_refs = {
+            parent_path: [("oneOf", child_path)],
+            child_path: [],
+        }
+        variant_needs = {parent_path: {"create", "update"}}
+
+        preprocess_schemas.propagate_needs_transitive(
+            variant_needs, schema_refs, schemas
+        )
+
+        self.assertEqual(variant_needs[child_path], {"create", "update"})
+
     def test_main_preprocesses_schema_tree_end_to_end(self) -> None:
         """The full pipeline normalizes schemas and writes linked variants."""
         with tempfile.TemporaryDirectory() as temp_dir: