From 8ab6b6ef9541d8950ad6d585026f657e9ef52fd2 Mon Sep 17 00:00:00 2001
From: Bilal Al-Shahwany <bilal@split.io>
Date: Tue, 22 Aug 2023 12:57:51 -0700
Subject: [PATCH 1/2] Update split model

---
 splitio/models/splits.py    | 22 +++++++++++++++++-----
 tests/models/test_splits.py |  5 ++++-
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/splitio/models/splits.py b/splitio/models/splits.py
index 5e0ab394..9cf0eba0 100644
--- a/splitio/models/splits.py
+++ b/splitio/models/splits.py
@@ -7,7 +7,7 @@
 
 SplitView = namedtuple(
     'SplitView',
-    ['name', 'traffic_type', 'killed', 'treatments', 'change_number', 'configs']
+    ['name', 'traffic_type', 'killed', 'treatments', 'change_number', 'configs', 'sets']
 )
 
 
@@ -41,7 +41,8 @@ def __init__(  # pylint: disable=too-many-arguments
             algo=None,
             traffic_allocation=None,
             traffic_allocation_seed=None,
-            configurations=None
+            configurations=None,
+            sets=[]
     ):
         """
         Class constructor.
@@ -62,6 +63,8 @@ def __init__(  # pylint: disable=too-many-arguments
         :type traffic_allocation: int
         :pram traffic_allocation_seed: Seed used to hash traffic allocation.
         :type traffic_allocation_seed: int
+        :pram sets: list of flag sets
+        :type sets: list
         """
         self._name = name
         self._seed = seed
@@ -90,6 +93,7 @@ def __init__(  # pylint: disable=too-many-arguments
             self._algo = HashAlgorithm.LEGACY
 
         self._configurations = configurations
+        self._sets = sets
 
     @property
     def name(self):
@@ -146,6 +150,11 @@ def traffic_allocation_seed(self):
         """Return the traffic allocation seed of the split."""
         return self._traffic_allocation_seed
 
+    @property
+    def sets(self):
+        """Return the flag sets of the split."""
+        return self._sets
+
     def get_configurations_for(self, treatment):
         """Return the mapping of treatments to configurations."""
         return self._configurations.get(treatment) if self._configurations else None
@@ -173,7 +182,8 @@ def to_json(self):
             'defaultTreatment': self.default_treatment,
             'algo': self.algo.value,
             'conditions': [c.to_json() for c in self.conditions],
-            'configurations': self._configurations
+            'configurations': self._configurations,
+            'sets': self._sets
         }
 
     def to_split_view(self):
@@ -189,7 +199,8 @@ def to_split_view(self):
             self.killed,
             list(set(part.treatment for cond in self.conditions for part in cond.partitions)),
             self.change_number,
-            self._configurations if self._configurations is not None else {}
+            self._configurations if self._configurations is not None else {},
+            self._sets
         )
 
     def local_kill(self, default_treatment, change_number):
@@ -238,5 +249,6 @@ def from_raw(raw_split):
         raw_split.get('algo'),
         traffic_allocation=raw_split.get('trafficAllocation'),
         traffic_allocation_seed=raw_split.get('trafficAllocationSeed'),
-        configurations=raw_split.get('configurations')
+        configurations=raw_split.get('configurations'),
+        sets=raw_split['sets']
     )
diff --git a/tests/models/test_splits.py b/tests/models/test_splits.py
index 847448b0..da289ad0 100644
--- a/tests/models/test_splits.py
+++ b/tests/models/test_splits.py
@@ -60,6 +60,7 @@ class SplitTests(object):
         'configurations': {
             'on': '{"color": "blue", "size": 13}'
         },
+        'sets': ['set1', 'set2']
     }
 
     def test_from_raw(self):
@@ -79,6 +80,7 @@ def test_from_raw(self):
         assert len(parsed.conditions) == 2
         assert parsed.get_configurations_for('on') == '{"color": "blue", "size": 13}'
         assert parsed._configurations == {'on': '{"color": "blue", "size": 13}'}
+        assert parsed.sets == ['set1', 'set2']
 
     def test_get_segment_names(self, mocker):
         """Test fetching segment names."""
@@ -89,7 +91,6 @@ def test_get_segment_names(self, mocker):
         split1 = splits.Split( 'some_split', 123, False, 'off', 'user', 'ACTIVE', 123, [cond1, cond2])
         assert split1.get_segment_names() == ['segment%d' % i for i in range(1, 5)]
 
-
     def test_to_json(self):
         """Test json serialization."""
         as_json = splits.from_raw(self.raw).to_json()
@@ -105,6 +106,7 @@ def test_to_json(self):
         assert as_json['defaultTreatment'] == 'off'
         assert as_json['algo'] == 2
         assert len(as_json['conditions']) == 2
+        assert as_json['sets'] == ['set1', 'set2']
 
     def test_to_split_view(self):
         """Test SplitView creation."""
@@ -115,3 +117,4 @@ def test_to_split_view(self):
         assert as_split_view.killed == self.raw['killed']
         assert as_split_view.traffic_type == self.raw['trafficTypeName']
         assert set(as_split_view.treatments) == set(['on', 'off'])
+        assert as_split_view.sets == self.raw['sets']

From 43df7f4da2d682bdab40fa2c03daa0003637fd64 Mon Sep 17 00:00:00 2001
From: Bilal Al-Shahwany <bilal@split.io>
Date: Tue, 22 Aug 2023 13:29:12 -0700
Subject: [PATCH 2/2] polish

---
 splitio/models/splits.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/splitio/models/splits.py b/splitio/models/splits.py
index 9cf0eba0..31f0fd0b 100644
--- a/splitio/models/splits.py
+++ b/splitio/models/splits.py
@@ -250,5 +250,5 @@ def from_raw(raw_split):
         traffic_allocation=raw_split.get('trafficAllocation'),
         traffic_allocation_seed=raw_split.get('trafficAllocationSeed'),
         configurations=raw_split.get('configurations'),
-        sets=raw_split['sets']
+        sets=raw_split.get('sets')
     )