FF-3390 fix(python): fix typing information for Configuration.__init__

python-sdk/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "eppo_py"
-version = "4.0.1"
+version = "4.0.2"
 edition = "2021"
 publish = false
 

python-sdk/python/eppo_client/_eppo_client.pyi

@@ -6,7 +6,9 @@ def init(config: ClientConfig) -> EppoClient: ...
 def get_instance() -> EppoClient: ...
 
 class Configuration:
-    def __init__(self, flags_configuration: bytes) -> None: ...
+    def __init__(
+        self, *, flags_configuration: bytes, bandits_configuration: bytes | None = None
+    ) -> None: ...
     def get_flags_configuration(self) -> bytes: ...
     def get_flag_keys(self) -> Set[str]: ...
     def get_bandit_keys(self) -> Set[str]: ...

python-sdk/tests/test_configuration.py

@@ -5,12 +5,104 @@
 
 from .util import init
 
+FLAGS_CONFIG = json.dumps(
+    {
+        "createdAt": "2024-09-09T10:18:15.988Z",
+        "environment": {"name": "test"},
+        "flags": {},
+    }
+).encode("utf-8")
+
+FLAGS_CONFIG_WITH_BANDITS = json.dumps(
+    {
+        "createdAt": "2024-09-09T10:18:15.988Z",
+        "environment": {"name": "test"},
+        "flags": {},
+        "bandits": {
+            "banner_bandit": [
+                {
+                    "key": "banner_bandit",
+                    "flagKey": "banner_bandit_flag",
+                    "variationKey": "banner_bandit",
+                    "variationValue": "banner_bandit",
+                },
+                {
+                    "key": "banner_bandit",
+                    "flagKey": "banner_bandit_flag_uk_only",
+                    "variationKey": "banner_bandit",
+                    "variationValue": "banner_bandit",
+                },
+            ],
+            "car_bandit": [
+                {
+                    "key": "car_bandit",
+                    "flagKey": "car_bandit_flag",
+                    "variationKey": "car_bandit",
+                    "variationValue": "car_bandit",
+                }
+            ],
+        },
+    }
+).encode("utf-8")
+
+BANDITS_MODEL_CONFIG = json.dumps(
+    {
+        "updatedAt": "2023-09-13T04:52:06.462Z",
+        "environment": {"name": "Test"},
+        "bandits": {
+            "car_bandit": {
+                "banditKey": "car_bandit",
+                "modelName": "falcon",
+                "updatedAt": "2023-09-13T04:52:06.462Z",
+                "modelVersion": "v456",
+                "modelData": {
+                    "gamma": 1.0,
+                    "defaultActionScore": 5.0,
+                    "actionProbabilityFloor": 0.2,
+                    "coefficients": {
+                        "toyota": {
+                            "actionKey": "toyota",
+                            "intercept": 1.0,
+                            "actionNumericCoefficients": [
+                                {
+                                    "attributeKey": "speed",
+                                    "coefficient": 1,
+                                    "missingValueCoefficient": 0.0,
+                                }
+                            ],
+                            "actionCategoricalCoefficients": [],
+                            "subjectNumericCoefficients": [],
+                            "subjectCategoricalCoefficients": [],
+                        }
+                    },
+                },
+            }
+        },
+    }
+).encode("utf-8")
+
 
 class TestConfiguration:
     def test_init_valid(self):
-        Configuration(
-            flags_configuration=b'{"createdAt":"2024-09-09T10:18:15.988Z","environment":{"name":"test"},"flags":{}}'
+        Configuration(flags_configuration=FLAGS_CONFIG)
+
+    def test_bandit_configuration_without_models(self):
+        config = Configuration(flags_configuration=FLAGS_CONFIG_WITH_BANDITS)
+
+        # Call get_bandit_keys and check the output
+        bandit_keys = config.get_bandit_keys()
+        assert isinstance(bandit_keys, set)
+        assert len(bandit_keys) == 0
+
+    def test_bandit_model_configuration(self):
+        config = Configuration(
+            flags_configuration=FLAGS_CONFIG_WITH_BANDITS,
+            bandits_configuration=BANDITS_MODEL_CONFIG,
         )
+        bandit_keys = config.get_bandit_keys()
+        assert isinstance(bandit_keys, set)
+        # `car_bandit` is the only bandit in the model config.
+        assert bandit_keys == {"car_bandit"}
 
     def test_init_invalid_json(self):
         """Input is not valid JSON string."""