Skip to content

Commit b2ee907

Browse files
authored
add custom id lookup in override (#274)
1 parent 168ffb7 commit b2ee907

File tree

2 files changed

+123
-15
lines changed

2 files changed

+123
-15
lines changed

statsig/evaluator.py

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,8 @@ def get_client_initialize_response(
102102
return None
103103

104104
return ClientInitializeResponseFormatter \
105-
.get_formatted_response(self.__eval_config, user, self._spec_store, self, hash, client_sdk_key, include_local_override)
105+
.get_formatted_response(self.__eval_config, user, self._spec_store, self, hash, client_sdk_key,
106+
include_local_override)
106107

107108
def _create_evaluation_details(self, reason: EvaluationReason):
108109
if reason == EvaluationReason.uninitialized:
@@ -118,7 +119,7 @@ def __lookup_gate_override(self, user, gate):
118119

119120
eval_details = self._create_evaluation_details(
120121
EvaluationReason.local_override)
121-
override = gate_overrides.get(user.user_id)
122+
override = self.__lookup_override(gate_overrides, user)
122123
if override is not None:
123124
return _ConfigEvaluation(boolean_value=override, rule_id="override",
124125
evaluation_details=eval_details)
@@ -130,7 +131,6 @@ def __lookup_gate_override(self, user, gate):
130131

131132
return None
132133

133-
134134
def lookup_gate_override(self, user, gate):
135135
return self.__lookup_gate_override(user, gate)
136136

@@ -141,7 +141,7 @@ def __lookup_config_override(self, user, config):
141141

142142
eval_details = self._create_evaluation_details(
143143
EvaluationReason.local_override)
144-
override = config_overrides.get(user.user_id)
144+
override = self.__lookup_override(config_overrides, user)
145145
if override is not None:
146146
return _ConfigEvaluation(json_value=override, rule_id="override",
147147
evaluation_details=eval_details)
@@ -152,7 +152,6 @@ def __lookup_config_override(self, user, config):
152152
json_value=all_override, rule_id="override", evaluation_details=eval_details)
153153
return None
154154

155-
156155
def lookup_config_override(self, user, config):
157156
return self.__lookup_config_override(user, config)
158157

@@ -163,7 +162,7 @@ def __lookup_layer_override(self, user, config):
163162

164163
eval_details = self._create_evaluation_details(
165164
EvaluationReason.local_override)
166-
override = layer_overrides.get(user.user_id)
165+
override = self.__lookup_override(layer_overrides, user)
167166
if override is not None:
168167
return _ConfigEvaluation(json_value=override, rule_id="override",
169168
evaluation_details=eval_details)
@@ -174,14 +173,23 @@ def __lookup_layer_override(self, user, config):
174173
json_value=all_override, rule_id="override", evaluation_details=eval_details)
175174
return None
176175

176+
def __lookup_override(self, config_overrides, user):
177+
override = config_overrides.get(user.user_id)
178+
if override is None and user.custom_ids is not None:
179+
for id_name in user.custom_ids:
180+
override = config_overrides.get(user.custom_ids[id_name])
181+
if override is not None:
182+
break
183+
return override
184+
177185
def unsupported_or_unrecognized(self, config_name):
178186
if config_name in self._spec_store.unsupported_configs:
179187
return _ConfigEvaluation(
180188
evaluation_details=self._create_evaluation_details(
181189
EvaluationReason.unsupported))
182190
return _ConfigEvaluation(
183-
evaluation_details=self._create_evaluation_details(
184-
EvaluationReason.unrecognized))
191+
evaluation_details=self._create_evaluation_details(
192+
EvaluationReason.unrecognized))
185193

186194
def check_gate(self, user, gate, end_result=None, is_nested=False):
187195
override = self.__lookup_gate_override(user, gate)

tests/test_local_mocks.py

Lines changed: 107 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ def test_override_gate(self):
3434

3535
user_one = StatsigUser("123", email="testuser@statsig.com")
3636
user_two = StatsigUser("456", email="test@statsig.com")
37+
user_custom = StatsigUser("789", custom_ids={"statsigId": "abc"})
3738

3839
self.assertEqual(
3940
server.check_gate(user_one, "any_gate"),
@@ -47,12 +48,25 @@ def test_override_gate(self):
4748
True
4849
)
4950

51+
self.assertEqual(
52+
server.check_gate(user_custom, "any_gate"),
53+
False
54+
)
55+
56+
server.override_gate("any_gate", True, "abc")
57+
58+
self.assertEqual(
59+
server.check_gate(user_custom, "any_gate"),
60+
True
61+
)
62+
5063
self.assertEqual(
5164
server.check_gate(user_two, "any_gate"),
5265
False
5366
)
5467

5568
server.override_gate("any_gate", False, "123")
69+
server.override_gate("any_gate", False, "abc")
5670
server.override_gate("any_gate", True, "456")
5771

5872
self.assertEqual(
@@ -79,22 +93,37 @@ def test_override_gate(self):
7993
server.check_gate(StatsigUser("4123980"), "any_gate"),
8094
True
8195
)
96+
self.assertEqual(
97+
server.check_gate(user_custom, "any_gate"),
98+
False
99+
)
82100

83101
# remove user override first
84102
server.remove_gate_override("any_gate", "123")
85-
# user one should now use global override
103+
server.remove_gate_override("any_gate", "abc")
104+
105+
# user one and user custom should now use global override
86106
self.assertEqual(
87107
server.check_gate(user_one, "any_gate"),
88108
True
89109
)
90110

111+
self.assertEqual(
112+
server.check_gate(user_custom, "any_gate"),
113+
True
114+
)
115+
91116
# remove global override
92117
server.remove_gate_override("any_gate")
93-
# user one should now have no override
118+
# user one and user custom should now have no override
94119
self.assertEqual(
95120
server.check_gate(user_one, "any_gate"),
96121
False
97122
)
123+
self.assertEqual(
124+
server.check_gate(user_custom, "any_gate"),
125+
False
126+
)
98127

99128
# remove all overrides
100129
server.remove_all_overrides()
@@ -113,6 +142,7 @@ def test_override_all(self):
113142

114143
user_one = StatsigUser("123", email="testuser@statsig.com")
115144
user_two = StatsigUser("456", email="test@statsig.com")
145+
user_custom = StatsigUser("789", custom_ids={"statsigId": "abc"})
116146

117147
self.assertEqual(
118148
server.check_gate(user_one, "any_gate"),
@@ -124,6 +154,11 @@ def test_override_all(self):
124154
False
125155
)
126156

157+
self.assertEqual(
158+
server.check_gate(user_custom, "any_gate"),
159+
False
160+
)
161+
127162
server.override_gate("any_gate", True)
128163

129164
self.assertEqual(
@@ -136,6 +171,11 @@ def test_override_all(self):
136171
True
137172
)
138173

174+
self.assertEqual(
175+
server.check_gate(user_custom, "any_gate"),
176+
True
177+
)
178+
139179
server.override_experiment("my_experiment", {"test": False})
140180

141181
self.assertEqual(
@@ -148,13 +188,19 @@ def test_override_all(self):
148188
{"test": False}
149189
)
150190

191+
self.assertEqual(
192+
server.get_experiment(user_custom, "my_experiment").get_value(),
193+
{"test": False}
194+
)
195+
151196
def test_override_config(self):
152197
options = StatsigOptions(local_mode=True, disable_diagnostics=True)
153198
server = StatsigServer()
154199
server.initialize("secret-key", options)
155200

156201
user_one = StatsigUser("123", email="testuser@statsig.com")
157202
user_two = StatsigUser("456", email="test@statsig.com")
203+
user_custom = StatsigUser("789", custom_ids={"statsigId": "abc"})
158204

159205
self.assertEqual(
160206
server.get_config(user_one, "config").get_value(),
@@ -166,20 +212,32 @@ def test_override_config(self):
166212
{}
167213
)
168214

215+
self.assertEqual(
216+
server.get_config(user_custom, "config").get_value(),
217+
{}
218+
)
219+
169220
override = {"test": 123}
170221
server.override_config("config", override, "123")
222+
server.override_config("config", override, "abc")
171223

172224
self.assertEqual(
173225
server.get_config(user_one, "config").get_value(),
174226
override
175227
)
176228

229+
self.assertEqual(
230+
server.get_config(user_custom, "config").get_value(),
231+
override
232+
)
233+
177234
self.assertEqual(
178235
server.get_config(user_two, "config").get_value(),
179236
{}
180237
)
181238

182239
server.override_experiment("config", {}, "123")
240+
server.override_experiment("config", {}, "def")
183241
new_override = {"abc": "def"}
184242
server.override_experiment("config", new_override, "456")
185243

@@ -188,6 +246,11 @@ def test_override_config(self):
188246
{}
189247
)
190248

249+
self.assertEqual(
250+
server.get_config(user_custom, "config").get_value(),
251+
override
252+
)
253+
191254
self.assertEqual(
192255
server.get_config(user_two, "config").get_value(),
193256
new_override
@@ -206,17 +269,30 @@ def test_override_config(self):
206269
new_override
207270
)
208271
self.assertEqual(
209-
server.get_config(StatsigUser("anyuser"), "config").get_value(),
272+
server.get_config(user_custom, "config").get_value(),
273+
override
274+
)
275+
self.assertEqual(
276+
server.get_config(StatsigUser("789"), "config").get_value(),
277+
new_override_2
278+
)
279+
self.assertEqual(
280+
server.get_config(StatsigUser("000", custom_ids={"statsigId": "ghi"}), "config").get_value(),
210281
new_override_2
211282
)
212283

213284
# remove user override first
214285
server.remove_config_override("config", "123")
215-
# user one should now use global override
286+
server.remove_config_override("config", "abc")
287+
# user one and user custom should now use global override
216288
self.assertEqual(
217289
server.get_config(user_one, "config").get_value(),
218290
new_override_2,
219291
)
292+
self.assertEqual(
293+
server.get_config(user_custom, "config").get_value(),
294+
new_override_2,
295+
)
220296

221297
# remove global override
222298
server.remove_config_override("config")
@@ -243,6 +319,7 @@ def test_override_layer(self):
243319

244320
user_one = StatsigUser("123", email="testuser@statsig.com")
245321
user_two = StatsigUser("456", email="test@statsig.com")
322+
user_custom = StatsigUser("789", custom_ids={"statsigId": "abc"})
246323

247324
# no override
248325
self.assertEqual(
@@ -253,9 +330,14 @@ def test_override_layer(self):
253330
server.get_layer(user_two, "a_layer").get("a_string", "fallback"),
254331
"fallback"
255332
)
333+
self.assertEqual(
334+
server.get_layer(user_custom, "a_layer").get("a_string", "fallback"),
335+
"fallback"
336+
)
256337

257338
# add override
258339
server.override_layer("a_layer", {"a_string": "override_str"}, "123")
340+
server.override_layer("a_layer", {"a_string": "override_str"}, "abc")
259341
self.assertEqual(
260342
server.get_layer(user_one, "a_layer").get("a_string", "fallback"),
261343
"override_str"
@@ -264,14 +346,23 @@ def test_override_layer(self):
264346
server.get_layer(user_two, "a_layer").get("a_string", "fallback"),
265347
"fallback"
266348
)
349+
self.assertEqual(
350+
server.get_layer(user_custom, "a_layer").get("a_string", "fallback"),
351+
"override_str"
352+
)
267353

268-
# change user 1, add to user 2
354+
# change user 1 and user custom, add to user 2
269355
server.override_layer("a_layer", {"a_bool": True}, "123")
356+
server.override_layer("a_layer", {"a_bool": True}, "abc")
270357
server.override_layer("a_layer", {"a_string": "override_str"}, "456")
271358
self.assertEqual(
272359
server.get_layer(user_one, "a_layer").get("a_string", "fallback"),
273360
"fallback"
274361
)
362+
self.assertEqual(
363+
server.get_layer(user_custom, "a_layer").get("a_string", "fallback"),
364+
"fallback"
365+
)
275366
self.assertEqual(
276367
server.get_layer(user_two, "a_layer").get("a_string", "fallback"),
277368
"override_str"
@@ -286,19 +377,28 @@ def test_override_layer(self):
286377

287378
# remove user override first
288379
server.remove_layer_override("a_layer", "123")
289-
# user one should now use global override
380+
server.remove_layer_override("a_layer", "abc")
381+
# user one and user custom should now use global override
290382
self.assertEqual(
291383
server.get_layer(user_one, "a_layer").get("a_string", "fallback"),
292384
"global_override"
293385
)
386+
self.assertEqual(
387+
server.get_layer(user_custom, "a_layer").get("a_string", "fallback"),
388+
"global_override"
389+
)
294390

295391
# remove global override
296392
server.remove_layer_override("a_layer")
297-
# user one should now have no override
393+
# user one and user custom should now have no override
298394
self.assertEqual(
299395
server.get_layer(user_one, "a_layer").get("a_string", "fallback"),
300396
"fallback"
301397
)
398+
self.assertEqual(
399+
server.get_layer(user_custom, "a_layer").get("a_string", "fallback"),
400+
"fallback"
401+
)
302402

303403
# remove all overrides
304404
server.remove_all_overrides()

0 commit comments

Comments
 (0)