From 19334294eb2f665afb55deb8fd882ce47c78bb17 Mon Sep 17 00:00:00 2001
From: Simon Dale <33907262+simondale00@users.noreply.github.com>
Date: Thu, 9 May 2024 11:00:33 -0400
Subject: [PATCH] [EN-20691] Added a getAssignmentVariation function (#10)

---
 src/main/java/com/eppo/sdk/EppoClient.java | 52 ++++++++++++++++++++--
 1 file changed, 49 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/eppo/sdk/EppoClient.java b/src/main/java/com/eppo/sdk/EppoClient.java
index 62ed5c3..de9fc34 100644
--- a/src/main/java/com/eppo/sdk/EppoClient.java
+++ b/src/main/java/com/eppo/sdk/EppoClient.java
@@ -46,6 +46,51 @@ protected Optional<EppoValue> getAssignmentValue(
             String flagKey,
             EppoAttributes subjectAttributes,
             Map<String, EppoAttributes> actionsWithAttributes
+    ) {
+        Optional<Variation> assignedVariation = getAssignmentVariation(
+                subjectKey,
+                flagKey,
+                subjectAttributes,
+                actionsWithAttributes
+        );
+
+        if (assignedVariation.isPresent()) {
+            return Optional.of(assignedVariation.get().getTypedValue());
+        }
+
+        return Optional.empty();
+    }
+
+    /**
+     * Returns the assigned variation.
+     *
+     * @param subjectKey
+     * @param flagKey
+     * @param subjectAttributes
+     * @return
+     */
+    public Optional<Variation> getAssignmentVariation(
+            String subjectKey,
+            String flagKey,
+            EppoAttributes subjectAttributes
+    ) {
+        return getAssignmentVariation(subjectKey, flagKey, subjectAttributes, null);
+    }
+
+    /**
+     * Returns the assigned variation.
+     *
+     * @param subjectKey
+     * @param flagKey
+     * @param subjectAttributes
+     * @param actionsWithAttributes
+     * @return
+     */
+    protected Optional<Variation> getAssignmentVariation(
+            String subjectKey,
+            String flagKey,
+            EppoAttributes subjectAttributes,
+            Map<String, EppoAttributes> actionsWithAttributes
     ) {
         // Validate Input Values
         InputValidator.validateNotBlank(subjectKey, "Invalid argument: subjectKey cannot be blank");
@@ -68,10 +113,11 @@ protected Optional<EppoValue> getAssignmentValue(
 
         if (algorithmType == AlgorithmType.OVERRIDE) {
             // Assigned variation was from an override; return its value without logging
-            return assignmentValue;
+            return Optional.of(assignedVariation);
         } else if (algorithmType == AlgorithmType.CONTEXTUAL_BANDIT) {
             // Assigned variation is a bandit; need to use the bandit to determine its value
-            assignmentValue = this.determineAndLogBanditAction(assignmentResult, actionsWithAttributes);
+            Optional<EppoValue> banditValue = this.determineAndLogBanditAction(assignmentResult, actionsWithAttributes);
+            assignedVariation.setTypedValue(banditValue.orElse(null));
         }
 
         // Log the assignment
@@ -88,7 +134,7 @@ protected Optional<EppoValue> getAssignmentValue(
             log.warn("Error logging assignment", e);
         }
 
-        return assignmentValue;
+        return Optional.of(assignedVariation);
     }
 
     private VariationAssignmentResult getAssignedVariation(String flagKey, String subjectKey, EppoAttributes subjectAttributes) {