Fix typing issues with Pydantic

Author: atainter

tests/test_fga.py

@@ -107,9 +107,8 @@ def setup(self, sync_http_client_for_test):
         self.http_client = sync_http_client_for_test
         self.fga = FGA(http_client=self.http_client)
 
-    @pytest.fixture
-    def mock_check_warning_response(self):
-        return {
+    def test_check_with_warning(self, mock_http_client_with_response):
+        mock_response = {
             "result": "authorized",
             "is_implicit": True,
             "warnings": [
@@ -120,10 +119,23 @@ def mock_check_warning_response(self):
                 }
             ],
         }
+        mock_http_client_with_response(self.http_client, mock_response, 200)
 
-    @pytest.fixture
-    def mock_query_warning_response(self):
-        return {
+        response = self.fga.check(
+            op="any_of",
+            checks=[
+                WarrantCheckInput(
+                    resource_type="schedule",
+                    resource_id="schedule-A1",
+                    relation="viewer",
+                    subject=SubjectInput(resource_type="user", resource_id="user-A"),
+                )
+            ],
+        )
+        assert response.dict(exclude_none=True) == mock_response
+
+    def test_query_with_warning(self, mock_http_client_with_response):
+        mock_response = {
             "object": "list",
             "data": [
                 {
@@ -149,12 +161,28 @@ def mock_query_warning_response(self):
             ],
         }
 
-    def test_check_with_warning(
-        self, mock_check_warning_response, mock_http_client_with_response
-    ):
-        mock_http_client_with_response(
-            self.http_client, mock_check_warning_response, 200
+        mock_http_client_with_response(self.http_client, mock_response, 200)
+
+        response = self.fga.query(
+            q="select member of type user for permission:view-docs",
+            order="asc",
+            warrant_token="warrant_token",
         )
+        assert response.dict(exclude_none=True) == mock_response
+
+    def test_check_with_generic_warning(self, mock_http_client_with_response):
+        mock_response = {
+            "result": "authorized",
+            "is_implicit": True,
+            "warnings": [
+                {
+                    "code": "generic",
+                    "message": "Generic warning",
+                }
+            ],
+        }
+
+        mock_http_client_with_response(self.http_client, mock_response, 200)
 
         response = self.fga.check(
             op="any_of",
@@ -167,21 +195,7 @@ def test_check_with_warning(
                 )
             ],
         )
-        assert response.dict(exclude_none=True) == mock_check_warning_response
-
-    def test_query_with_warning(
-        self, mock_query_warning_response, mock_http_client_with_response
-    ):
-        mock_http_client_with_response(
-            self.http_client, mock_query_warning_response, 200
-        )
-
-        response = self.fga.query(
-            q="select member of type user for permission:view-docs",
-            order="asc",
-            warrant_token="warrant_token",
-        )
-        assert response.dict(exclude_none=True) == mock_query_warning_response
+        assert response.dict(exclude_none=True) == mock_response
 
 
 class TestFGA:

workos/fga.py

@@ -10,7 +10,8 @@
     WarrantWrite,
     WarrantWriteOperation,
     WriteWarrantResponse,
-    WarrantQueryResult, FGAWarning,
+    WarrantQueryResult,
+    FGAWarning,
 )
 from workos.types.fga.list_filters import (
     AuthorizationResourceListFilters,
@@ -46,9 +47,9 @@
 WarrantListResource = WorkOSListResource[Warrant, WarrantListFilters, ListMetadata]
 
 
-class WarrantQueryListResource(WorkOSListResource[
-                                   WarrantQueryResult, WarrantQueryListFilters, ListMetadata
-                               ]):
+class WarrantQueryListResource(
+    WorkOSListResource[WarrantQueryResult, WarrantQueryListFilters, ListMetadata]
+):
     warnings: Optional[Sequence[FGAWarning]] = None
 
 

workos/types/fga/__init__.py

@@ -2,3 +2,4 @@
 from .authorization_resource_types import *
 from .authorization_resources import *
 from .warrant import *
+from .warnings import *

workos/types/fga/warnings.py

@@ -1,6 +1,7 @@
-from typing import Sequence, Union, Literal, Annotated
+from typing import Sequence, Annotated, Union, Any, Dict
 
-from pydantic import Field
+from pydantic import BeforeValidator
+from pydantic_core.core_schema import ValidationInfo
 
 from workos.types.workos_model import WorkOSModel
 
@@ -11,11 +12,18 @@ class FGABaseWarning(WorkOSModel):
 
 
 class MissingContextKeysWarning(FGABaseWarning):
-    code: Literal["missing_context_keys"]
     keys: Sequence[str]
 
 
+def fga_warning_dispatch_validator(
+    value: Dict[str, Any], info: ValidationInfo
+) -> FGABaseWarning:
+    if value.get("code") == "missing_context_keys":
+        return MissingContextKeysWarning.model_validate(value)
+    return FGABaseWarning.model_validate(value)
+
+
 FGAWarning = Annotated[
     Union[MissingContextKeysWarning, FGABaseWarning],
-    Field(discriminator='type')
+    BeforeValidator(fga_warning_dispatch_validator),
 ]