Touch up tests for new version of pydantic

PrefectHQ · Sep 9, 2024 · ce77eac · ce77eac
1 parent 3db5fd8
commit ce77eac
Show file tree

Hide file tree

Showing 13 changed files with 65 additions and 71 deletions.
diff --git a/src/integrations/prefect-kubernetes/tests/test_worker.py b/src/integrations/prefect-kubernetes/tests/test_worker.py
@@ -449,7 +449,7 @@ def enable_store_api_key_in_secret(monkeypatch):
                     "cluster_config": {
                         "title": "Cluster Config",
                         "description": "The Kubernetes cluster config to use for job creation.",
-                        "allOf": [{"$ref": "#/definitions/KubernetesClusterConfig"}],
+                        "$ref": "#/definitions/KubernetesClusterConfig",
                     },
                 },
                 "definitions": {
@@ -870,7 +870,7 @@ def enable_store_api_key_in_secret(monkeypatch):
                     "cluster_config": {
                         "title": "Cluster Config",
                         "description": "The Kubernetes cluster config to use for job creation.",
-                        "allOf": [{"$ref": "#/definitions/KubernetesClusterConfig"}],
+                        "$ref": "#/definitions/KubernetesClusterConfig",
                     },
                     "memory": {
                         "title": "Memory",

diff --git a/src/prefect/infrastructure/base.py b/src/prefect/infrastructure/base.py
@@ -104,7 +104,7 @@ async def generate_work_pool_base_job_template(self):
                         "description": (
                             "The infrastructure block to use for job creation."
                         ),
-                        "allOf": [{"$ref": f"#/definitions/{self.__class__.__name__}"}],
+                        "$ref": f"#/definitions/{self.__class__.__name__}",
                         "default": {
                             "$ref": {"block_document_id": str(self._block_document_id)}
                         },

diff --git a/tests/fixtures/collections_registry.py b/tests/fixtures/collections_registry.py
@@ -173,7 +173,7 @@ def k8s_default_base_job_template():
                     "description": (
                         "The Kubernetes cluster config to use for job creation."
                     ),
-                    "allOf": [{"$ref": "#/definitions/KubernetesClusterConfig"}],
+                    "$ref": "#/definitions/KubernetesClusterConfig",
                 },
             },
             "definitions": {

diff --git a/tests/infrastructure/provisioners/test_cloud_run.py b/tests/infrastructure/provisioners/test_cloud_run.py
@@ -124,7 +124,7 @@
                     "The GCP Credentials used to initiate the Cloud Run Job."
                     " `service_account_info` is required."
                 ),
-                "allOf": [{"$ref": "#/definitions/GcpCredentials"}],
+                "$ref": "#/definitions/GcpCredentials",
             },
             "image": {
                 "title": "Image Name",

diff --git a/tests/infrastructure/provisioners/test_cloud_run_v2.py b/tests/infrastructure/provisioners/test_cloud_run_v2.py
@@ -87,7 +87,7 @@
             "credentials": {
                 "title": "GCP Credentials",
                 "description": "The GCP Credentials used to connect to Cloud Run. If not provided credentials will be inferred from the local environment.",
-                "allOf": [{"$ref": "#/definitions/GcpCredentials"}],
+                "$ref": "#/definitions/GcpCredentials",
             },
             "region": {
                 "title": "Region",

diff --git a/tests/infrastructure/provisioners/test_container_instance.py b/tests/infrastructure/provisioners/test_container_instance.py
@@ -272,9 +272,7 @@ def default_base_job_template():
                 "aci_credentials": {
                     "title": "Aci Credentials",
                     "description": "The credentials to use to authenticate with Azure.",
-                    "allOf": [
-                        {"$ref": "#/definitions/AzureContainerInstanceCredentials"}
-                    ],
+                    "$ref": "#/definitions/AzureContainerInstanceCredentials",
                 },
                 "stream_output": {
                     "title": "Stream Output",

diff --git a/tests/infrastructure/test_base.py b/tests/infrastructure/test_base.py
@@ -320,7 +320,7 @@ async def test_generate_work_pool_base_job_template():
                 "block": {
                     "title": "Block",
                     "description": "The infrastructure block to use for job creation.",
-                    "allOf": [{"$ref": "#/definitions/MockInfrastructure"}],
+                    "$ref": "#/definitions/MockInfrastructure",
                     "default": {
                         "$ref": {"block_document_id": str(block._block_document_id)}
                     },
@@ -358,7 +358,7 @@ async def test_publish_as_work_pool(
                 "block": {
                     "title": "Block",
                     "description": "The infrastructure block to use for job creation.",
-                    "allOf": [{"$ref": "#/definitions/MockInfrastructure"}],
+                    "$ref": "#/definitions/MockInfrastructure",
                     "default": {
                         "$ref": {"block_document_id": str(block._block_document_id)}
                     },

diff --git a/tests/server/api/test_deployments.py b/tests/server/api/test_deployments.py
@@ -1069,7 +1069,7 @@ async def test_create_deployment_parameter_enforcement_allows_partial_parameters
                         "position": 1,
                     },
                     "person": {
-                        "allOf": [{"$ref": "#/definitions/Person"}],
+                        "$ref": "#/definitions/Person",
                         "position": 0,
                     },
                 },
@@ -1441,7 +1441,7 @@ async def test_update_deployment_parameter_enforcement_allows_partial_parameters
                         "position": 1,
                     },
                     "person": {
-                        "allOf": [{"$ref": "#/definitions/Person"}],
+                        "$ref": "#/definitions/Person",
                         "position": 0,
                     },
                 },
@@ -2821,7 +2821,7 @@ async def test_create_flow_run_from_deployment_parameter_enforcement_rejects_par
                         "position": 1,
                     },
                     "person": {
-                        "allOf": [{"$ref": "#/definitions/Person"}],
+                        "$ref": "#/definitions/Person",
                         "position": 0,
                     },
                 },

diff --git a/tests/server/api/test_infra_overrides.py b/tests/server/api/test_infra_overrides.py
@@ -407,7 +407,7 @@ async def test_base_job_template_default_references_to_blocks(
                     "required": ["k8s_credentials"],
                     "properties": {
                         "k8s_credentials": {
-                            "allOf": [{"$ref": "#/definitions/k8s_credentials"}],
+                            "$ref": "#/definitions/k8s_credentials",
                             "title": "K8s Credentials",
                             "default": {
                                 "$ref": {
@@ -673,7 +673,7 @@ async def test_base_job_template_default_references_to_blocks(
                     "required": ["k8s_credentials"],
                     "properties": {
                         "k8s_credentials": {
-                            "allOf": [{"$ref": "#/definitions/k8s_credentials"}],
+                            "$ref": "#/definitions/k8s_credentials",
                             "title": "K8s Credentials",
                             "default": {
                                 "$ref": {

diff --git a/tests/server/api/ui/test_schemas.py b/tests/server/api/ui/test_schemas.py
@@ -102,7 +102,7 @@ async def test_circular_schema_reference(
                         "param": {
                             "title": "param",
                             "position": 0,
-                            "allOf": [{"$ref": "#/definitions/City"}],
+                            "$ref": "#/definitions/City",
                         }
                     },
                     "required": ["param"],
@@ -118,7 +118,7 @@ async def test_circular_schema_reference(
                             },
                             "required": ["population", "name"],
                             # City definition references itself here
-                            "allOf": [{"$ref": "#/definitions/City"}],
+                            "$ref": "#/definitions/City",
                         }
                     },
                 },

diff --git a/tests/utilities/schema_tools/test_validation.py b/tests/utilities/schema_tools/test_validation.py
@@ -61,7 +61,7 @@ async def test_circular_schema_ref():
             "param": {
                 "title": "param",
                 "position": 0,
-                "allOf": [{"$ref": "#/definitions/City"}],
+                "$ref": "#/definitions/City",
             }
         },
         "required": ["param"],
@@ -74,7 +74,7 @@ async def test_circular_schema_ref():
                 },
                 "required": ["population", "name"],
                 # City definition references itself here
-                "allOf": [{"$ref": "#/definitions/City"}],
+                "$ref": "#/definitions/City",
             }
         },
     }
@@ -357,7 +357,7 @@ def schema(self) -> dict:
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/City"}],
+                    "$ref": "#/definitions/City",
                 }
             },
             "required": ["param"],
@@ -415,7 +415,7 @@ def schema(self) -> dict:
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/City"}],
+                    "$ref": "#/definitions/City",
                 }
             },
             "required": ["param"],
@@ -634,7 +634,7 @@ def schema(self) -> dict:
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/City"}],
+                    "$ref": "#/definitions/City",
                 }
             },
             "required": ["param"],
@@ -1145,7 +1145,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1176,7 +1176,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1206,7 +1206,7 @@ def test_pydantic_v1_model_required_int_or_none(self):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1233,7 +1233,7 @@ def test_pydantic_v1_model_required_int_or_none(self):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1259,7 +1259,7 @@ def test_pydantic_v1_model_optional_int_or_none_default_none(self):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1287,7 +1287,7 @@ def test_pydantic_v1_model_optional_int_or_none_default_none(self):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1317,7 +1317,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1473,7 +1473,7 @@ def test_pydantic_v2_model_required_int(self):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -1503,7 +1503,7 @@ def test_pydantic_v2_model_optional_int(self):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -1536,7 +1536,7 @@ def test_pydantic_v2_model_optional_int_default_none(self):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -1569,7 +1569,7 @@ def test_pydantic_v2_model_required_int_or_none(self):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -1602,7 +1602,7 @@ def test_pydantic_v2_model_optional_int_or_none_default_none(self):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -1635,7 +1635,7 @@ def test_pydantic_v2_model_optional_int_or_none_default_int(self):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -1789,7 +1789,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1819,7 +1819,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1853,7 +1853,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1886,7 +1886,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1922,7 +1922,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -1954,7 +1954,7 @@ class MyModel(BaseModel):
                 "param": {
                     "title": "param",
                     "position": 0,
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                 }
             },
             "required": ["param"],
@@ -2061,7 +2061,7 @@ class MyModel(BaseModel):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -2098,7 +2098,7 @@ class MyModel(BaseModel):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }
@@ -2137,7 +2137,7 @@ class MyModel(BaseModel):
             "type": "object",
             "properties": {
                 "param": {
-                    "allOf": [{"$ref": "#/definitions/MyModel"}],
+                    "$ref": "#/definitions/MyModel",
                     "position": 0,
                     "title": "param",
                 }