fix: restore trait value coercion to string

Flagsmith · Jul 4, 2023 · 6973191 · 6973191
1 parent 298c32b
commit 6973191
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/flag_engine/identities/traits/models.py b/flag_engine/identities/traits/models.py
@@ -1,4 +1,7 @@
-from pydantic import BaseModel
+from typing import Any, get_args
+
+from pydantic import BaseModel, validator
+
 
 from flag_engine.identities.traits.types import TraitValue
 
@@ -7,5 +10,11 @@ class TraitModel(BaseModel):
     trait_key: str
     trait_value: TraitValue = ...
 
+    @validator("trait_value", pre=True)
+    def convert_trait_value(cls, value: Any) -> TraitValue:
+        if isinstance(value, get_args(TraitValue)):
+            return value
+        return str(value)
+
     class Config:
         smart_union: bool = True
diff --git a/tests/unit/identities/test_identities_models.py b/tests/unit/identities/test_identities_models.py
@@ -186,6 +186,9 @@ def test_get_hash_key_with_use_identity_composite_key_for_hashing_disabled(ident
         (False, False),
         (0.0, 0.0),
         (0, 0),
+        (None, None),
+        ([], "[]"),
+        (["SUPERADMIN"], "['SUPERADMIN']"),
     ],
 )
 def test_trait_model__deserialize__expected_trait_value(