improve logic, test coverage

Author: khvn26

src/main/java/com/flagsmith/mappers/EngineMappers.java

@@ -18,6 +18,7 @@
 import com.flagsmith.models.FeatureMetadata;
 import com.flagsmith.models.Flag;
 import com.flagsmith.models.SegmentMetadata;
+import com.flagsmith.models.TraitModel;
 import com.flagsmith.models.environments.EnvironmentModel;
 import com.flagsmith.models.features.FeatureModel;
 import com.flagsmith.models.features.FeatureSegmentModel;
@@ -28,6 +29,7 @@
 import com.flagsmith.models.segments.SegmentConditionModel;
 import com.flagsmith.models.segments.SegmentModel;
 import com.flagsmith.models.segments.SegmentRuleModel;
+import com.flagsmith.utils.ModelUtils;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -79,22 +81,16 @@ public static EvaluationContext mapContextAndIdentityDataToContext(
     // Create identity context
     IdentityContext identityContext = new IdentityContext()
         .withIdentifier(identifier)
-        .withKey(context.getEnvironment().getKey() + "_" + identifier)
-        .withTraits(new Traits());
+        .withKey(context.getEnvironment().getKey() + "_" + identifier);
 
     // Map traits if provided
     if (traits != null && !traits.isEmpty()) {
-      for (Map.Entry<String, Object> entry : traits.entrySet()) {
-        Object traitValue = entry.getValue();
-        // Handle TraitConfig-like objects (maps with "value" key)
-        if (traitValue instanceof Map) {
-          Map<?, ?> traitMap = (Map<?, ?>) traitValue;
-          if (traitMap.containsKey("value")) {
-            traitValue = traitMap.get("value");
-          }
-        }
-        identityContext.getTraits().setAdditionalProperty(entry.getKey(), traitValue);
+      Traits identityTraits = new Traits();
+      for (TraitModel traitModel : ModelUtils.getTraitModelsFromTraitMap(traits)) {
+        identityTraits.setAdditionalProperty(
+            traitModel.getTraitKey(), traitModel.getTraitValue());
       }
+      identityContext.setTraits(identityTraits);
     }
 
     // Create new evaluation context with identity

src/test/java/com/flagsmith/flagengine/unit/mappers/EngineMappersTest.java

@@ -0,0 +1,52 @@
+package com.flagsmith.flagengine.unit.mappers;
+
+import com.flagsmith.mappers.EngineMappers;
+import com.flagsmith.models.TraitConfig;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.util.Map;
+import java.util.stream.Stream;
+
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import com.flagsmith.FlagsmithTestHelper;
+import com.flagsmith.flagengine.EvaluationContext;
+import com.flagsmith.flagengine.Traits;
+
+public class EngineMappersTest {
+  private static Stream<Arguments> expectedTraitMaps() {
+    return Stream.of(
+        Arguments.argumentSet(
+          "no transiency data",
+          Map.of("test", 1)
+        ),
+        Arguments.argumentSet(
+          "with transiency data",
+          Map.of("test", TraitConfig.fromObject(1))
+        )
+    );
+  }
+
+  @ParameterizedTest
+  @MethodSource("expectedTraitMaps")
+  public void testMapContextAndIdentityDataToContext_returnsExpectedContext(
+    Map<String, Object> expectedTraitMap
+  ) {
+    // Arrange
+    final String identifier = "test-identifier";
+    final EvaluationContext context = FlagsmithTestHelper.evaluationContext();
+    final Traits expectedTraits = new Traits().withAdditionalProperty("test", 1);
+
+    // Act
+    final EvaluationContext mappedContext = EngineMappers.mapContextAndIdentityDataToContext(
+      context, identifier, expectedTraitMap);
+
+    // Assert
+    assertEquals(
+      expectedTraits.getAdditionalProperties(),
+      mappedContext.getIdentity().getTraits().getAdditionalProperties());
+  }
+}