diff --git a/infrahub_sdk/checks.py b/infrahub_sdk/checks.py index 09d1b30c..a873893c 100644 --- a/infrahub_sdk/checks.py +++ b/infrahub_sdk/checks.py @@ -11,7 +11,7 @@ from pydantic import BaseModel, Field from . import InfrahubClient -from .exceptions import InfrahubCheckNotFoundError +from .exceptions import InfrahubCheckNotFoundError, UninitializedError if TYPE_CHECKING: from pathlib import Path @@ -41,6 +41,7 @@ def __init__( output: Optional[str] = None, initializer: Optional[InfrahubCheckInitializer] = None, params: Optional[dict] = None, + client: Optional[InfrahubClient] = None, ): self.git: Optional[Repo] = None self.initializer = initializer or InfrahubCheckInitializer() @@ -55,7 +56,7 @@ def __init__( self.root_directory = root_directory or os.getcwd() - self.client: InfrahubClient + self._client = client if not self.name: self.name = self.__class__.__name__ @@ -66,6 +67,17 @@ def __init__( def __str__(self) -> str: return self.__class__.__name__ + @property + def client(self) -> InfrahubClient: + if self._client: + return self._client + + raise UninitializedError(message="This check has not been initialized with a client") + + @client.setter + def client(self, value: InfrahubClient) -> None: + self._client = value + @classmethod async def init(cls, client: Optional[InfrahubClient] = None, *args: Any, **kwargs: Any) -> InfrahubCheck: """Async init method, If an existing InfrahubClient client hasn't been provided, one will be created automatically."""