Include test cases

datamodel_code_generator/model/pydantic_v2/root_model.py

@@ -15,6 +15,7 @@ def __init__(
     ) -> None:
         # Remove custom_base_class for Pydantic V2 models; behaviour is different from Pydantic V1 as it will not
         # be treated as a root model. custom_base_class cannot both implement BaseModel and RootModel!
-        kwargs.pop('custom_base_class')
+        if 'custom_base_class' in kwargs:
+            kwargs.pop('custom_base_class')
 
         super().__init__(**kwargs)

tests/model/pydantic_v2/test_root_model.py

@@ -0,0 +1,67 @@
+from datamodel_code_generator.model.pydantic_v2.root_model import RootModel
+from datamodel_code_generator.model import DataModelFieldBase
+from datamodel_code_generator.reference import Reference
+from datamodel_code_generator.types import DataType
+
+
+def test_root_model():
+    root_model = RootModel(
+        fields=[
+            DataModelFieldBase(
+                name='a',
+                data_type=DataType(type='str'),
+                default='abc',
+                required=False,
+            )
+        ],
+        reference=Reference(name='TestRootModel', path='test_root_model'),
+    )
+
+    assert root_model.name == 'TestRootModel'
+    assert root_model.fields == [
+        DataModelFieldBase(
+            name='a',
+            data_type=DataType(type='str'),
+            default='abc',
+            required=False,
+            parent=root_model,
+        )
+    ]
+
+    assert root_model.base_class == 'RootModel'
+    assert root_model.custom_base_class is None
+    assert root_model.render() == (
+        'class TestRootModel(RootModel):\n' "    root: Optional[str] = 'abc'"
+    )
+
+
+def test_root_model_custom_base_class():
+    root_model = RootModel(
+        custom_base_class='test.Test',
+        fields=[
+            DataModelFieldBase(
+                name='a',
+                data_type=DataType(type='str'),
+                default='abc',
+                required=False,
+            )
+        ],
+        reference=Reference(name='TestRootModel', path='test_root_model'),
+    )
+
+    assert root_model.name == 'TestRootModel'
+    assert root_model.fields == [
+        DataModelFieldBase(
+            name='a',
+            data_type=DataType(type='str'),
+            default='abc',
+            required=False,
+            parent=root_model,
+        )
+    ]
+
+    assert root_model.base_class == 'RootModel'
+    assert root_model.custom_base_class is None
+    assert root_model.render() == (
+        'class TestRootModel(RootModel):\n' "    root: Optional[str] = 'abc'"
+    )