Raise a proper branch not found exception for branches that doesn't exist

changelog/286.fixed.md

@@ -0,0 +1 @@
+Raise a proper branch not found error when requesting a node or schema for a branch that doesn't exist.

infrahub_sdk/schema/__init__.py

@@ -13,6 +13,7 @@
 from typing_extensions import TypeAlias
 
 from ..exceptions import (
+    BranchNotFoundError,
     InvalidResponseError,
     JsonDecodeError,
     SchemaNotFoundError,
@@ -167,6 +168,25 @@
 
         raise ValueError("schema must be a protocol or a string")
 
+    @staticmethod
+    def _parse_schema_response(response: httpx.Response, branch: str) -> MutableMapping[str, Any]:
+        if response.status_code == 400:
+            raise BranchNotFoundError(
+                identifier=branch, message=f"The requested branch was not found on the server [{branch}]"
+            )
+        response.raise_for_status()
+
+        try:
+            data: MutableMapping[str, Any] = response.json()
+        except json.decoder.JSONDecodeError as exc:
+            raise JsonDecodeError(
+                message=f"Invalid Schema response received from the server at {response.url}: {response.text} [{response.status_code}] ",
+                content=response.text,
+                url=str(response.url),
+            ) from exc
+
+        return data
+
 
 class InfrahubSchema(InfrahubSchemaBase):
     def __init__(self, client: InfrahubClient):
@@ -420,16 +440,8 @@
         url = f"{self.client.address}/api/schema?{query_params}"
 
         response = await self.client._get(url=url, timeout=timeout)
-        response.raise_for_status()
 
-        try:
-            data: MutableMapping[str, Any] = response.json()
-        except json.decoder.JSONDecodeError as exc:
-            raise JsonDecodeError(
-                message=f"Invalid Schema response received from the server at {response.url}: {response.text} [{response.status_code}] ",
-                content=response.text,
-                url=response.url,
-            ) from exc
+        data = self._parse_schema_response(response=response, branch=branch)
 
         nodes: MutableMapping[str, MainSchemaTypesAPI] = {}
         for node_schema in data.get("nodes", []):
@@ -648,9 +660,7 @@
         query_params = urlencode(url_parts)
         url = f"{self.client.address}/api/schema?{query_params}"
         response = self.client._get(url=url, timeout=timeout)
-        response.raise_for_status()
-
-        data: MutableMapping[str, Any] = response.json()
+        data = self._parse_schema_response(response=response, branch=branch)
 
         nodes: MutableMapping[str, MainSchemaTypesAPI] = {}
         for node_schema in data.get("nodes", []):

tests/integration/test_schema.py

@@ -1,15 +1,27 @@
-# import pytest
+import pytest
+
+from infrahub_sdk import InfrahubClient
+from infrahub_sdk.exceptions import BranchNotFoundError
+from infrahub_sdk.testing.docker import TestInfrahubDockerClient
+
+
 # from infrahub.core.schema import core_models
 # from infrahub.server import app
 #
-# from infrahub_sdk import Config, InfrahubClient
 # from infrahub_sdk.schema import NodeSchemaAPI
 #
 # from .conftest import InfrahubTestClient
 #
 #
 #
-#
+class TestInfrahubSchema(TestInfrahubDockerClient):
+    async def test_query_schema_for_branch_not_found(self, client: InfrahubClient):
+        with pytest.raises(BranchNotFoundError) as exc:
+            await client.all(kind="BuiltinTag", branch="I-do-not-exist")
+
+        assert str(exc.value) == "The requested branch was not found on the server [I-do-not-exist]"
+
+
 # class TestInfrahubSchema:
 #     @pytest.fixture(scope="class")
 #     async def client(self):