@@ -121,6 +121,37 @@ def test_is_valid_traffic_type(self, mocker):
121121 assert storage .is_valid_traffic_type ('user' ) is False
122122 assert storage .is_valid_traffic_type ('account' ) is False
123123
124+ def test_traffic_type_inc_dec_logic (self , mocker ):
125+ """Test that adding/removing split, handles traffic types correctly."""
126+ storage = InMemorySplitStorage ()
127+
128+ split1 = mocker .Mock ()
129+ name1_prop = mocker .PropertyMock ()
130+ name1_prop .return_value = 'split1'
131+ type(split1 ).name = name1_prop
132+
133+ split2 = mocker .Mock ()
134+ name2_prop = mocker .PropertyMock ()
135+ name2_prop .return_value = 'split1'
136+ type(split2 ).name = name2_prop
137+
138+ tt_user = mocker .PropertyMock ()
139+ tt_user .return_value = 'user'
140+
141+ tt_account = mocker .PropertyMock ()
142+ tt_account .return_value = 'account'
143+
144+ type(split1 ).traffic_type_name = tt_user
145+ type(split2 ).traffic_type_name = tt_account
146+
147+ storage .put (split1 )
148+ assert storage .is_valid_traffic_type ('user' ) is True
149+ assert storage .is_valid_traffic_type ('account' ) is False
150+
151+ storage .put (split2 )
152+ assert storage .is_valid_traffic_type ('user' ) is False
153+ assert storage .is_valid_traffic_type ('account' ) is True
154+
124155
125156class InMemorySegmentStorageTests (object ):
126157 """In memory segment storage tests."""
0 commit comments