From 4cbb87fbab82ae7b6b64819c6734e5c447381978 Mon Sep 17 00:00:00 2001
From: Justin Abrahms <justin@abrah.ms>
Date: Sun, 7 Aug 2022 21:08:45 -0700
Subject: [PATCH] Support chaining evaluation context additions.

Refs #29
---
 .../openfeature/javasdk/EvaluationContext.java    | 15 ++++++++++-----
 .../dev/openfeature/javasdk/EvalContextTest.java  |  9 +++++++++
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/main/java/dev/openfeature/javasdk/EvaluationContext.java b/src/main/java/dev/openfeature/javasdk/EvaluationContext.java
index 6dd9a597..7f22a7f3 100644
--- a/src/main/java/dev/openfeature/javasdk/EvaluationContext.java
+++ b/src/main/java/dev/openfeature/javasdk/EvaluationContext.java
@@ -21,8 +21,9 @@ public class EvaluationContext {
 
     // TODO Not sure if I should have sneakythrows or checked exceptions here..
     @SneakyThrows
-    public <T> void addStructureAttribute(String key, T value) {
+    public <T> EvaluationContext addStructureAttribute(String key, T value) {
         attributes.put(key, new Pair<>(FlagValueType.OBJECT, value));
+        return this;
     }
 
     @SneakyThrows
@@ -34,8 +35,9 @@ public class EvaluationContext {
         return getAttributesByType(FlagValueType.OBJECT);
     }
 
-    public void addStringAttribute(String key, String value) {
+    public EvaluationContext addStringAttribute(String key, String value) {
         attributes.put(key, new Pair<>(FlagValueType.STRING, value));
+        return this;
     }
 
     public String getStringAttribute(String key) {
@@ -68,8 +70,9 @@ public class EvaluationContext {
         return getAttributesByType(FlagValueType.STRING);
     }
 
-    public void addIntegerAttribute(String key, Integer value) {
+    public EvaluationContext addIntegerAttribute(String key, Integer value) {
         attributes.put(key, new Pair<>(FlagValueType.INTEGER, value));
+        return this;
     }
 
     public Integer getIntegerAttribute(String key) {
@@ -84,16 +87,18 @@ public class EvaluationContext {
         return getAttributeByType(key, FlagValueType.BOOLEAN);
     }
 
-    public void addBooleanAttribute(String key, Boolean b) {
+    public EvaluationContext addBooleanAttribute(String key, Boolean b) {
         attributes.put(key, new Pair<>(FlagValueType.BOOLEAN, b));
+        return this;
     }
 
     public Map<String, Boolean> getBooleanAttributes() {
         return getAttributesByType(FlagValueType.BOOLEAN);
     }
 
-    public void addDatetimeAttribute(String key, ZonedDateTime value) {
+    public EvaluationContext addDatetimeAttribute(String key, ZonedDateTime value) {
         attributes.put(key, new Pair<>(FlagValueType.STRING, value.format(DateTimeFormatter.ISO_ZONED_DATE_TIME)));
+        return this;
     }
 
     public ZonedDateTime getDatetimeAttribute(String key) {
diff --git a/src/test/java/dev/openfeature/javasdk/EvalContextTest.java b/src/test/java/dev/openfeature/javasdk/EvalContextTest.java
index 4dc12ef1..cb29270c 100644
--- a/src/test/java/dev/openfeature/javasdk/EvalContextTest.java
+++ b/src/test/java/dev/openfeature/javasdk/EvalContextTest.java
@@ -109,4 +109,13 @@ public class EvalContextTest {
         assertEquals(null, ec.getStringAttribute("key"));
         assertEquals(3, ec.getIntegerAttribute("key"));
     }
+
+    @Test void can_chain_attribute_addition() {
+        EvaluationContext ec = new EvaluationContext();
+        EvaluationContext out = ec.addStructureAttribute("str", "test")
+                .addIntegerAttribute("int", 4)
+                .addBooleanAttribute("bool", false)
+                .addStructureAttribute("str", new Node<Integer>());
+        assertEquals(EvaluationContext.class, out.getClass());
+    }
 }