Feeback from review

Author: dgarros

changelog/+schema-set-cache.added.md

@@ -1 +1 @@
-Add method schema.set_cache() to populate the cache manually (primarily for unit testing)
+Add method `client.schema.set_cache()` to populate the cache manually (primarily for unit testing)

changelog/+schema-timeout.deprecated.md

@@ -1 +1 @@
-The 'timeout' parameter  while fetching the schema has been deprecated. the default_timeout will be used instead.
+The 'timeout' parameter while creating a node or fetching the schema has been deprecated. the default_timeout will be used instead.

infrahub_sdk/schema/__init__.py

@@ -110,10 +110,19 @@ def validate_data_against_schema(self, schema: MainSchemaTypesAPI, data: dict) -
                     message=f"{key} is not a valid value for {identifier}",
                 )
 
-    def set_cache(self, schema: dict[str, Any] | BranchSchema, branch: str | None = None) -> None:
+    def set_cache(self, schema: dict[str, Any] | SchemaRootAPI | BranchSchema, branch: str | None = None) -> None:
+        """
+        Set the cache manually (primarily for unit testing)
+
+        Args:
+            schema: The schema to set the cache as provided by the /api/schema endpoint either in dict or SchemaRootAPI format
+            branch: The name of the branch to set the cache for.
+        """
         branch = branch or self.client.default_branch
 
-        if isinstance(schema, dict):
+        if isinstance(schema, SchemaRootAPI):
+            schema = BranchSchema.from_schema_root_api(data=schema)
+        elif isinstance(schema, dict):
             schema = BranchSchema.from_api_response(data=schema)
 
         self.cache[branch] = schema

infrahub_sdk/schema/main.py

@@ -345,7 +345,7 @@ def to_schema_dict(self) -> dict[str, Any]:
 class SchemaRootAPI(BaseModel):
     model_config = ConfigDict(use_enum_values=True)
 
-    version: str
+    main: str | None = None
     generics: list[GenericSchemaAPI] = Field(default_factory=list)
     nodes: list[NodeSchemaAPI] = Field(default_factory=list)
     profiles: list[ProfileSchemaAPI] = Field(default_factory=list)
@@ -360,23 +360,29 @@ class BranchSchema(BaseModel):
 
     @classmethod
     def from_api_response(cls, data: MutableMapping[str, Any]) -> Self:
+        """
+        Convert an API response from /api/schema into a BranchSchema object.
+        """
+        return cls.from_schema_root_api(data=SchemaRootAPI(**data))
+
+    @classmethod
+    def from_schema_root_api(cls, data: SchemaRootAPI) -> Self:
+        """
+        Convert a SchemaRootAPI object to a BranchSchema object.
+        """
         nodes: MutableMapping[str, GenericSchemaAPI | NodeSchemaAPI | ProfileSchemaAPI | TemplateSchemaAPI] = {}
-        for node_schema in data.get("nodes", []):
-            node = NodeSchemaAPI(**node_schema)
+        for node in data.nodes:
             nodes[node.kind] = node
 
-        for generic_schema in data.get("generics", []):
-            generic = GenericSchemaAPI(**generic_schema)
+        for generic in data.generics:
             nodes[generic.kind] = generic
 
-        for profile_schema in data.get("profiles", []):
-            profile = ProfileSchemaAPI(**profile_schema)
+        for profile in data.profiles:
             nodes[profile.kind] = profile
 
-        for template_schema in data.get("templates", []):
-            template = TemplateSchemaAPI(**template_schema)
+        for template in data.templates:
             nodes[template.kind] = template
 
-        schema_hash = data.get("main", "")
+        schema_hash = data.main or ""
 
         return cls(hash=schema_hash, nodes=nodes)