fix coverage

Author: khvn26

flag_engine/engine.py

@@ -107,7 +107,7 @@ def _get_identity_feature_states_dict(
         identity=identity,
     )
     if override_traits:
-        if typing.TYPE_CHECKING:
+        if typing.TYPE_CHECKING:  # pragma: no cover
             assert context["identity"]
         context["identity"].setdefault("traits", {}).update(
             {trait.trait_key: trait.trait_value for trait in override_traits}

flag_engine/segments/evaluator.py

@@ -105,9 +105,11 @@ def context_matches_condition(
 
 
 def _get_trait(context: EvaluationContext, trait_key: str) -> TraitValue:
-    if identity_context := context["identity"]:
-        return identity_context["traits"][trait_key]
-    return None
+    return (
+        identity_context["traits"][trait_key]
+        if (identity_context := context["identity"])
+        else None
+    )
 
 
 def get_context_value(

tests/unit/segments/test_segments_evaluator.py

@@ -218,7 +218,7 @@ def test_context_in_segment(
     "segment_split_value, identity_hashed_percentage, expected_result",
     ((10, 1, True), (100, 50, True), (0, 1, False), (10, 20, False)),
 )
-def test_identity_in_segment_percentage_split(
+def test_context_in_segment_percentage_split(
     mocker: MockerFixture,
     context: EvaluationContext,
     segment_split_value: int,
@@ -246,6 +246,36 @@ def test_identity_in_segment_percentage_split(
     assert result == expected_result
 
 
+def test_context_in_segment_percentage_split__trait_value__calls_expected(
+    mocker: MockerFixture,
+    context: EvaluationContext,
+) -> None:
+    # Given
+    assert context["identity"] is not None
+    context["identity"]["traits"]["custom_trait"] = "custom_value"
+    percentage_split_condition = SegmentConditionModel(
+        operator=constants.PERCENTAGE_SPLIT,
+        value="10",
+        property_="custom_trait",
+    )
+    rule = SegmentRuleModel(
+        type=constants.ALL_RULE, conditions=[percentage_split_condition]
+    )
+    segment = SegmentModel(id=1, name="% split", rules=[rule])
+
+    mock_get_hashed_percentage = mocker.patch(
+        "flag_engine.segments.evaluator.get_hashed_percentage_for_object_ids"
+    )
+    mock_get_hashed_percentage.return_value = 1
+
+    # When
+    result = is_context_in_segment(context=context, segment=segment)
+
+    # Then
+    mock_get_hashed_percentage.assert_called_once_with([segment.id, "custom_value"])
+    assert result
+
+
 @pytest.mark.parametrize(
     "operator,  property_, expected_result",
     (