Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
Add `InExperiment` Udf (#203)
+53
-2
+1
CHANGELOG.md
+1
CHANGELOG.md
···
15
15
### 🎉 New features
16
16
- Add Postgres execution result store ([#171](https://github.com/roostorg/osprey/pull/171) by [@serendipty01](https://github.com/serendipty01))
17
17
- Add `StringSlice` UDF which extracts a substring by index range ([#189](https://github.com/roostorg/osprey/pull/189) by [@bealsbe](https://github.com/bealsbe))
18
+
- Add `InExperiment` UDF which checks if an entity is in an experiment ([#203](https://github.com/roostorg/osprey/pull/203) by [@bealsbe](https://github.com/bealsbe))
18
19
19
20
### 🐛 Bug fixes
20
21
- Default to selecting all for event stream ([#194](https://github.com/roostorg/osprey/pull/194) by [@chimosky](https://github.com/chimosky))
+16
osprey_worker/src/osprey/engine/stdlib/udfs/experiments.py
+16
osprey_worker/src/osprey/engine/stdlib/udfs/experiments.py
···
299
299
return arguments.extra_arguments[bucket]
300
300
301
301
302
+
class InExperimentArguments(ArgumentsBase):
303
+
experiment: ExperimentT
304
+
305
+
306
+
class InExperiment(UDFBase[InExperimentArguments, bool]):
307
+
"""
308
+
Returns True if the entity was assigned to any bucket in the experiment,
309
+
False if the entity fell outside all bucket ranges
310
+
"""
311
+
312
+
category = UdfCategories.ENGINE
313
+
314
+
def execute(self, execution_context: ExecutionContext, arguments: InExperimentArguments) -> bool:
315
+
return arguments.experiment.resolved_bucket is not NOT_IN_EXPERIMENT_BUCKET
316
+
317
+
302
318
ExperimentsBucketAssignment = Experiment.build_cls('experiment_bucket_assignment')
+34
-2
osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_experiments.py
+34
-2
osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_experiments.py
···
7
7
from osprey.engine.conftest import CheckFailureFunction, ExecuteFunction, RunValidationFunction
8
8
from osprey.engine.language_types.experiments import NOT_IN_EXPERIMENT_BUCKET, NOT_IN_EXPERIMENT_BUCKET_INDEX
9
9
from osprey.engine.stdlib.udfs.entity import Entity
10
-
from osprey.engine.stdlib.udfs.experiments import CONTROL_BUCKET, EXPERIMENT_GRANULARITY, Experiment, ExperimentWhen
10
+
from osprey.engine.stdlib.udfs.experiments import (
11
+
CONTROL_BUCKET,
12
+
EXPERIMENT_GRANULARITY,
13
+
Experiment,
14
+
ExperimentWhen,
15
+
InExperiment,
16
+
)
11
17
from osprey.engine.stdlib.udfs.rules import Rule
12
18
from osprey.engine.udf.registry import UDFRegistry
13
19
14
20
pytestmark: List[Callable[[Any], Any]] = [
15
-
pytest.mark.use_udf_registry(UDFRegistry.with_udfs(Entity, Rule, Experiment, ExperimentWhen)),
21
+
pytest.mark.use_udf_registry(UDFRegistry.with_udfs(Entity, Rule, Experiment, ExperimentWhen, InExperiment)),
16
22
pytest.mark.use_validators([ValidateCallKwargs, UniqueStoredNames]),
17
23
]
18
24
···
367
373
"""
368
374
data = execute(experiment)
369
375
assert data['B'] == expected_value
376
+
377
+
378
+
@mock.patch.object(Experiment, 'hash_mod')
379
+
def test_inexperiment_returns_true_when_in_experiment(hash_mod_mock: mock.MagicMock, execute: ExecuteFunction) -> None:
380
+
hash_mod_mock.return_value = 0
381
+
experiment = f"""
382
+
E1 = Entity(type='MyEntity', id='entity 1')
383
+
A = Experiment(entity=E1, buckets=['{CONTROL_BUCKET}', 'b'], bucket_sizes=[50.0, 50.0], version=1, revision=1)
384
+
B = InExperiment(experiment=A)
385
+
"""
386
+
data = execute(experiment)
387
+
assert data['B'] is True
388
+
389
+
390
+
@mock.patch.object(Experiment, 'hash_mod')
391
+
def test_inexperiment_returns_false_when_not_in_experiment(
392
+
hash_mod_mock: mock.MagicMock, execute: ExecuteFunction
393
+
) -> None:
394
+
hash_mod_mock.return_value = 9999
395
+
experiment = f"""
396
+
E1 = Entity(type='MyEntity', id='entity 1')
397
+
A = Experiment(entity=E1, buckets=['{CONTROL_BUCKET}', 'b'], bucket_sizes=[1.0, 1.0], version=1, revision=1)
398
+
B = InExperiment(experiment=A)
399
+
"""
400
+
data = execute(experiment)
401
+
assert data['B'] is False
370
402
371
403
372
404
@mock.patch.object(Experiment, 'hash_mod')
+2
osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py
+2
osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py
···
9
9
Experiment,
10
10
ExperimentsBucketAssignment,
11
11
ExperimentWhen,
12
+
InExperiment,
12
13
)
13
14
from osprey.engine.stdlib.udfs.extract_cookie import ExtractCookie
14
15
from osprey.engine.stdlib.udfs.get_action_name import GetActionName
···
82
83
Experiment,
83
84
ExperimentWhen,
84
85
ExperimentsBucketAssignment,
86
+
InExperiment,
85
87
ExtractCookie,
86
88
GetActionName,
87
89
HasLabel,