Set correct fields for RetryConfig

Author: FlorentinD

graphdatascience/arrow_client/authenticated_flight_client.py

@@ -56,8 +56,8 @@ def create(
                     FlightUnavailableError,
                     FlightInternalError,
                 ],
-                stop=StopConfig(after_delay=10, after_attempt=5),
-                wait=ExponentialWaitConfig(multiplier=1, min=1, max=10),
+                stop_config=StopConfig(after_delay=10, after_attempt=5),
+                wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
             )
 
         return AuthenticatedArrowClient(

graphdatascience/retry_utils/retry_config.py

@@ -21,7 +21,10 @@ class RetryConfig:
     retry: retry_base
 
 
-class RetryConfigV2(BaseModel):
+class RetryConfigV2(
+    BaseModel,
+    extra="forbid",
+):
     """Retry configuration which can be serialized/deserialized."""
 
     retryable_exceptions: list[type[BaseException]] | None = None
@@ -109,7 +112,7 @@ def tenacity_retry(self) -> retry_base | None:
         return retry_any(*retries)
 
 
-class StopConfig(BaseModel):
+class StopConfig(BaseModel, extra="forbid"):
     after_attempt: int | None = None
     after_delay: int | None = None
 
@@ -128,7 +131,7 @@ def to_tenacity(self) -> stop_base | None:
         return stop_any(*stops)
 
 
-class ExponentialWaitConfig(BaseModel):
+class ExponentialWaitConfig(BaseModel, extra="forbid"):
     multiplier: float = 1
     min: int = 0
     max: float | int = sys.maxsize / 2

graphdatascience/tests/unit/arrow_client/test_authenticated_flight_client.py

@@ -35,12 +35,10 @@ def auth_pair(self) -> tuple[str, str]:
     return MockAuthentication()
 
 
-def test_create_authenticated_arrow_client(
-    arrow_info: ArrowInfo, retry_config: RetryConfigV2, mock_auth: ArrowAuthentication
-) -> None:
-    client = AuthenticatedArrowClient.create(
-        arrow_info=arrow_info, auth=mock_auth, encrypted=True, retry_config=retry_config
-    )
+def test_create_authenticated_arrow_client(arrow_info: ArrowInfo, mock_auth: ArrowAuthentication) -> None:
+    client = AuthenticatedArrowClient.create(arrow_info=arrow_info, auth=mock_auth, encrypted=True)
+
+    assert client._retry_config.wait_config == ExponentialWaitConfig(multiplier=1, min=1, max=10)
     assert isinstance(client, AuthenticatedArrowClient)
     assert client.connection_info() == ConnectionInfo("localhost", 8491, encrypted=True)