diff --git a/growthbook.py b/growthbook.py index 4ea9e7a..b240dec 100644 --- a/growthbook.py +++ b/growthbook.py @@ -366,22 +366,6 @@ class Filter(TypedDict): hashVersion: int attribute: str -class StickyBucketServiceInterface: - def get_assignments(self, attributeName: str, attributeValue: str) -> Optional[Dict]: - pass - def save_assignments(self, doc: Dict) -> None: - pass - # By default, just loop through all attributes and call get_assignments - # Override this method in subclasses to perform a multi-query instead - def get_all_assignments(self, attributes: Dict[str, str]) -> Dict[str, Dict]: - docs = {} - for attributeName, attributeValue in attributes.items(): - doc = self.get_assignments(attributeName, attributeValue) - if doc: - key = f"{doc['attributeName']}||{doc['attributeValue']}" - docs[key] = doc - return docs - class Experiment(object): def __init__( self, @@ -777,6 +761,41 @@ def clear(self) -> None: self.cache.clear() +class AbstractStickyBucketService(ABC): + @abstractmethod + def get_assignments(self, attributeName: str, attributeValue: str) -> Optional[Dict]: + pass + + @abstractmethod + def save_assignments(self, doc: Dict) -> None: + pass + + def get_key(self, attributeName: str, attributeValue: str) -> str: + return f"{attributeName}||{attributeValue}" + + # By default, just loop through all attributes and call get_assignments + # Override this method in subclasses to perform a multi-query instead + def get_all_assignments(self, attributes: Dict[str, str]) -> Dict[str, Dict]: + docs = {} + for attributeName, attributeValue in attributes.items(): + doc = self.get_assignments(attributeName, attributeValue) + if doc: + docs[self.get_key(attributeName, attributeValue)] = doc + return docs + +class InMemoryStickyBucketService(AbstractStickyBucketService): + def __init__(self) -> None: + self.docs: Dict[str, Dict] = {} + + def get_assignments(self, attributeName: str, attributeValue: str) -> Optional[Dict]: + return self.docs.get(self.get_key(attributeName, attributeValue), None) + + def save_assignments(self, doc: Dict) -> None: + self.docs[self.get_key(doc["attributeName"], doc["attributeValue"])] = doc + + def destroy(self) -> None: + self.docs.clear() + class FeatureRepository(object): def __init__(self) -> None: self.cache: AbstractFeatureCache = InMemoryFeatureCache() @@ -870,9 +889,8 @@ def __init__( decryption_key: str = "", cache_ttl: int = 60, forced_variations: dict = {}, - sticky_bucket_service: StickyBucketServiceInterface = None, + sticky_bucket_service: AbstractStickyBucketService = None, sticky_bucket_identifier_attributes: List[str] = None, - sticky_bucket_assignment_docs: Dict[str, dict] = {}, # Deprecated args trackingCallback=None, qaMode: bool = False, @@ -891,10 +909,9 @@ def __init__( self._cache_ttl = cache_ttl self.sticky_bucket_identifier_attributes = sticky_bucket_identifier_attributes self.sticky_bucket_service = sticky_bucket_service - self.sticky_bucket_assignment_docs = sticky_bucket_assignment_docs - - if features: - self.setFeatures(features) + self.sticky_bucket_assignment_docs: dict = {} + self._using_derived_sticky_bucket_attributes = not sticky_bucket_identifier_attributes + self._sticky_bucket_attributes: Optional[dict] = None self._qaMode = qa_mode or qaMode self._trackingCallback = on_experiment_viewed or trackingCallback @@ -909,6 +926,9 @@ def __init__( self._assigned: Dict[str, Any] = {} self._subscriptions: Set[Any] = set() + if features: + self.setFeatures(features) + def load_features(self) -> None: if not self._client_key: raise ValueError("Must specify `client_key` to refresh features") @@ -933,6 +953,7 @@ def set_features(self, features: dict) -> None: rules=feature.get("rules", []), defaultValue=feature.get("defaultValue", None), ) + self.refresh_sticky_buckets() # @deprecated, use get_features def getFeatures(self) -> Dict[str, Feature]: @@ -947,6 +968,7 @@ def setAttributes(self, attributes: dict) -> None: def set_attributes(self, attributes: dict) -> None: self._attributes = attributes + self.refresh_sticky_buckets() # @deprecated, use get_attributes def getAttributes(self) -> dict: @@ -1275,6 +1297,9 @@ def _run(self, experiment: Experiment, featureId: Optional[str] = None) -> Resul found_sticky_bucket = sticky_bucket.get('variation', 0) >= 0 assigned = sticky_bucket.get('variation', 0) sticky_bucket_version_is_blocked = sticky_bucket.get('versionIsBlocked', False) + + if found_sticky_bucket: + logger.debug("Found sticky bucket for experiment %s, assigning sticky variation %s", experiment.key, assigned) # Some checks are not needed if we already have a sticky bucket if not found_sticky_bucket: @@ -1495,10 +1520,10 @@ def _getExperimentResult( stickyBucketUsed=stickyBucketUsed ) - def _derive_sticky_bucket_identifier_attributes(self, data: Dict[str, Feature]) -> List[str]: + def _derive_sticky_bucket_identifier_attributes(self) -> List[str]: attributes = set() - for key in data: - feature = data[key] + for key in self._features: + feature = self._features[key] for rule in feature.rules: if rule.variations: attributes.add(rule.hashAttribute or "id") @@ -1506,10 +1531,14 @@ def _derive_sticky_bucket_identifier_attributes(self, data: Dict[str, Feature]) attributes.add(rule.fallbackAttribute) return list(attributes) - def _get_sticky_bucket_attributes(self, data: Dict[str, Feature] = None) -> dict: + def _get_sticky_bucket_attributes(self) -> dict: attributes: Dict[str, str] = {} + if self._using_derived_sticky_bucket_attributes: + self.sticky_bucket_identifier_attributes = self._derive_sticky_bucket_identifier_attributes() + if not self.sticky_bucket_identifier_attributes: - self.sticky_bucket_identifier_attributes = self._derive_sticky_bucket_identifier_attributes(data or self._features) + return attributes + for attr in self.sticky_bucket_identifier_attributes: (_, hash_value) = self._getHashValue(attr) if hash_value: @@ -1563,15 +1592,17 @@ def _get_sticky_bucket_variation(self, experiment_key: str, bucket_version: int def _get_sticky_bucket_experiment_key(self, experiment_key: str, bucket_version: int = 0) -> str: return experiment_key + "__" + str(bucket_version) - def refresh_sticky_buckets(self, data: Dict[str, Feature] = None) -> None: + def refresh_sticky_buckets(self, force: bool = False) -> None: if not self.sticky_bucket_service: return - attributes = self._get_sticky_bucket_attributes(data) - self.sticky_bucket_assignment_docs = self.sticky_bucket_service.get_all_assignments(attributes) + attributes = self._get_sticky_bucket_attributes() + if not force and attributes == self._sticky_bucket_attributes: + logger.debug("Skipping refresh of sticky bucket assignments, no changes") + return - def get_sticky_bucket_assignment_docs(self) -> Dict[str, dict]: - return self.sticky_bucket_assignment_docs + self._sticky_bucket_attributes = attributes + self.sticky_bucket_assignment_docs = self.sticky_bucket_service.get_all_assignments(attributes) def _generate_sticky_bucket_assignment_doc(self, attribute_name: str, attribute_value: str, assignments: dict): key = attribute_name + "||" + attribute_value diff --git a/tests/cases.json b/tests/cases.json index 117e2ed..95edb08 100644 --- a/tests/cases.json +++ b/tests/cases.json @@ -4953,7 +4953,7 @@ { "attributes": { "id": "i123", - "anonymousId": "ses123", + "anonymousId": "123", "foo": "bar", "country": "USA" }, @@ -4965,7 +4965,7 @@ "key": "feature-exp", "seed": "feature-exp", "hashAttribute": "id", - "fallbackAttribute": "deviceId", + "fallbackAttribute": "anonymousId", "hashVersion": 2, "bucketVersion": 0, "condition": { "country": "USA" }, diff --git a/tests/test_growthbook.py b/tests/test_growthbook.py index 28031eb..ccb6b35 100644 --- a/tests/test_growthbook.py +++ b/tests/test_growthbook.py @@ -7,7 +7,7 @@ GrowthBook, Experiment, Feature, - StickyBucketServiceInterface, + InMemoryStickyBucketService, getBucketRanges, gbhash, chooseVariation, @@ -22,6 +22,7 @@ ) from time import time import pytest +from typing import Optional, Dict logger.setLevel("DEBUG") @@ -149,25 +150,31 @@ def test_stickyBucket(stickyBucket_data): _, ctx, key, expected_result, expected_docs = stickyBucket_data # Just use the interface directly, which passes and doesn't persist anywhere - ctx['sticky_bucket_service'] = StickyBucketServiceInterface() + service = InMemoryStickyBucketService() + ctx['sticky_bucket_service'] = service if 'stickyBucketIdentifierAttributes' in ctx: ctx['sticky_bucket_identifier_attributes'] = ctx['stickyBucketIdentifierAttributes'] ctx.pop('stickyBucketIdentifierAttributes') if 'stickyBucketAssignmentDocs' in ctx: - ctx['sticky_bucket_assignment_docs'] = ctx['stickyBucketAssignmentDocs'] + service.docs = ctx['stickyBucketAssignmentDocs'] ctx.pop('stickyBucketAssignmentDocs') gb = GrowthBook(**ctx) res = gb.eval_feature(key) + print(service.docs) + if not res.experimentResult: assert None == expected_result else: assert res.experimentResult.to_dict() == expected_result - assert gb.get_sticky_bucket_assignment_docs() == expected_docs + assert service.docs == expected_docs + + service.destroy() + gb.destroy() def getTrackingMock(gb: GrowthBook): @@ -740,3 +747,92 @@ def test_load_features(mocker): feature_repo.clear_cache() gb.destroy() + +def test_loose_unmarshalling(mocker): + m = mocker.patch.object(feature_repo, "_get") + m.return_value = MockHttpResp( + 200, json.dumps({"features": {"feature": {"defaultValue": 5, "rules": [{"force": 3, "unknown": "foo"}], "unknown": "foo"}}, "unknown": "foo"}) + ) + + gb = GrowthBook(api_host="https://cdn.growthbook.io", client_key="sdk-abc123") + + assert m.call_count == 0 + + gb.load_features() + m.assert_called_once_with("https://cdn.growthbook.io/api/features/sdk-abc123") + + assert gb.get_features()["feature"].to_dict() == {"defaultValue": 5, "rules": [{"force": 3, "hashVersion": 1}]} + + feature_repo.clear_cache() + gb.destroy() + +def test_sticky_bucket_service(mocker): + # Start forcing everyone to variation1 + features = { + "feature": { + "defaultValue": 5, + "rules": [{ + "key": "exp", + "variations": [0, 1], + "weights": [0, 1], + "meta": [ + {"key": "control"}, + {"key": "variation1"} + ] + }] + }, + } + + service = InMemoryStickyBucketService() + gb = GrowthBook( + sticky_bucket_service=service, + attributes={ + "id": "1" + }, + features=features + ) + + assert gb.get_feature_value("feature", -1) == 1 + assert service.get_assignments("id", "1") == { + "attributeName": "id", + "attributeValue": "1", + "assignments": { + "exp__0": "variation1" + } + } + + logger.debug("Change weights and ensure old user still gets variation") + features["feature"]["rules"][0]["weights"] = [1, 0] + gb.set_features(features) + assert gb.get_feature_value("feature", -1) == 1 + + logger.debug("New GrowthBook instance should also get variation") + gb2 = GrowthBook( + sticky_bucket_service=service, + attributes={ + "id": "1" + }, + features=features + ) + assert gb2.get_feature_value("feature", -1) == 1 + gb2.destroy() + + logger.debug("New users should get control") + gb.set_attributes({"id": "2"}) + assert gb.get_feature_value("feature", -1) == 0 + + logger.debug("Bumping bucketVersion, should reset sticky buckets") + gb.set_attributes({"id": "1"}) + features["feature"]["rules"][0]["bucketVersion"] = 1 + gb.set_features(features) + assert gb.get_feature_value("feature", -1) == 0 + + assert service.get_assignments("id", "1") == { + "attributeName": "id", + "attributeValue": "1", + "assignments": { + "exp__0": "variation1", + "exp__1": "control" + } + } + gb.destroy() \ No newline at end of file