diff --git a/data-model/src/main/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtils.java b/data-model/src/main/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtils.java index f4bcd4f..806969a 100644 --- a/data-model/src/main/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtils.java +++ b/data-model/src/main/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtils.java @@ -15,6 +15,10 @@ public static boolean containsAttribute(StructuredTrace trace, String attribute) && trace.getAttributes().getAttributeMap().containsKey(attribute); } + public static boolean getBooleanAttribute(StructuredTrace trace, String attribute) { + return Boolean.parseBoolean(getStringAttribute(trace, attribute)); + } + public static String getStringAttribute(StructuredTrace trace, String attribute) { if (trace.getAttributes() == null) { return null; diff --git a/data-model/src/test/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtilsTest.java b/data-model/src/test/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtilsTest.java index a3a91ff..f3def88 100644 --- a/data-model/src/test/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtilsTest.java +++ b/data-model/src/test/java/org/hypertrace/core/datamodel/shared/TraceAttributeUtilsTest.java @@ -1,7 +1,9 @@ package org.hypertrace.core.datamodel.shared; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -35,4 +37,38 @@ public void getListAttributeTest() { List.of("somevalue1", "somevalue2"), TraceAttributeUtils.getListAttribute(mockTrace, "someattribute")); } + + @Test + public void getStringAttributeTest() { + StructuredTrace mockTrace = mock(StructuredTrace.class); + assertNull(TraceAttributeUtils.getStringAttribute(mockTrace, "someattribute")); + + Map attributeMap = new HashMap<>(); + when(mockTrace.getAttributes()) + .thenReturn(Attributes.newBuilder().setAttributeMap(attributeMap).build()); + assertNull(TraceAttributeUtils.getStringAttribute(mockTrace, "someattribute")); + + attributeMap.put("someattribute", AttributeValue.newBuilder().setValue("somevalue").build()); + assertEquals("somevalue", TraceAttributeUtils.getStringAttribute(mockTrace, "someattribute")); + } + + @Test + public void getBooleanAttributeTest() { + StructuredTrace mockTrace = mock(StructuredTrace.class); + assertFalse(TraceAttributeUtils.getBooleanAttribute(mockTrace, "someattribute")); + + Map attributeMap = new HashMap<>(); + when(mockTrace.getAttributes()) + .thenReturn(Attributes.newBuilder().setAttributeMap(attributeMap).build()); + assertFalse(TraceAttributeUtils.getBooleanAttribute(mockTrace, "someattribute")); + + attributeMap.put("someattribute", AttributeValue.newBuilder().setValue("true").build()); + assertTrue(TraceAttributeUtils.getBooleanAttribute(mockTrace, "someattribute")); + + attributeMap.put("someOtherAttribute", AttributeValue.newBuilder().setValue("xyz").build()); + assertFalse(TraceAttributeUtils.getBooleanAttribute(mockTrace, "someOtherAttribute")); + + attributeMap.put("otherAttribute", AttributeValue.newBuilder().setValue("").build()); + assertFalse(TraceAttributeUtils.getBooleanAttribute(mockTrace, "otherAttribute")); + } }