From fbb958bd20f67554d2e433dccaf41f693e2ff787 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Tue, 11 Jul 2023 19:24:45 +0100
Subject: [PATCH] Fix Nan validation

---
 flag_engine/identities/traits/types.py          |  4 ++--
 tests/unit/identities/test_identities_models.py | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/flag_engine/identities/traits/types.py b/flag_engine/identities/traits/types.py
index f5de361f..92e16331 100644
--- a/flag_engine/identities/traits/types.py
+++ b/flag_engine/identities/traits/types.py
@@ -1,12 +1,12 @@
 from typing import Union
 
-from pydantic import constr
+from pydantic.types import confloat, constr
 
 from flag_engine.identities.traits.constants import TRAIT_STRING_VALUE_MAX_LENGTH
 
 TraitValue = Union[
     None,
-    float,
+    confloat(allow_inf_nan=False),
     int,
     bool,
     constr(max_length=TRAIT_STRING_VALUE_MAX_LENGTH),
diff --git a/tests/unit/identities/test_identities_models.py b/tests/unit/identities/test_identities_models.py
index 11a15dd9..426fb7cb 100644
--- a/tests/unit/identities/test_identities_models.py
+++ b/tests/unit/identities/test_identities_models.py
@@ -202,3 +202,18 @@ def test_trait_model__deserialize__expected_trait_value(
     # Then
     assert isinstance(result.trait_value, type(expected_result))
     assert result.trait_value == expected_result
+
+
+def test_identity_model__deserialize__handles_nan():
+    # When
+    result = IdentityModel.parse_obj(
+        {
+            "identifier": "identity",
+            "environment_api_key": "api-key",
+            "identity_traits": [{"trait_key": "nan_trait", "trait_value": "Nan"}],
+        }
+    )
+
+    # Then
+    assert result
+    assert result.identity_traits[0].trait_value == "Nan"