From eb2bb16091098e38203e50b18382cf74666603be Mon Sep 17 00:00:00 2001
From: Leo Romanovsky <leo@geteppo.com>
Date: Mon, 14 Oct 2024 21:11:21 -0700
Subject: [PATCH] Add new test case for initializating bandits config

---
 python-sdk/tests/test_configuration.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/python-sdk/tests/test_configuration.py b/python-sdk/tests/test_configuration.py
index 5d829770..3b437601 100644
--- a/python-sdk/tests/test_configuration.py
+++ b/python-sdk/tests/test_configuration.py
@@ -5,12 +5,20 @@
 
 from .util import init
 
+FLAGS_CONFIG = json.dumps({
+    "createdAt": "2024-09-09T10:18:15.988Z",
+    "environment": {"name": "test"},
+    "flags": {}
+}).encode('utf-8')
+
+BANDITS_CONFIG = json.dumps({
+    "bandit1": {"type": "some_type"},
+    "bandit2": {"type": "another_type"}
+}).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_init_invalid_json(self):
         """Input is not valid JSON string."""
@@ -21,17 +29,15 @@ def test_init_invalid_format(self):
         """flags is specified as array instead of object"""
         with pytest.raises(Exception):
             Configuration(
-                flags_configuration=b'{"createdAt":"2024-09-09T10:18:15.988Z","environment":{"name":"test"},"flags":[]}'
+                flags_configuration=FLAGS_CONFIG
             )
 
-
 @pytest.mark.rust_only
 def test_configuration_none():
     client = init("ufc", wait_for_init=False)
     configuration = client.get_configuration()
     assert configuration == None
 
-
 @pytest.mark.rust_only
 def test_configuration_some():
     client = init("ufc", wait_for_init=True)
@@ -43,3 +49,11 @@ def test_configuration_some():
     result = json.loads(flag_config)
     assert result["environment"] == {"name": "Test"}
     assert "numeric_flag" in result["flags"]
+
+def test_bandit_configuration():
+    # Initialize Configuration with both flags and bandits
+    config = Configuration(flags_configuration=FLAGS_CONFIG, bandits_configuration=BANDITS_CONFIG)
+
+    bandit_keys = config.get_bandit_keys()
+    assert isinstance(bandit_keys, set)
+    assert bandit_keys == {"bandit1", "bandit2"}