Repository: Unleash/unleash-client-java
Pull request #306 (2025-06-04 12:13:47)

chore: update to wasm engine
Moves the internals to use the new Yggdrasil WASM engine. The WASM engine collapses isEnabled/getVariant, impressionData and metrics into a one shot call, so updates here to conform to that new API


Patch:
diff --git a/pom.xml b/pom.xml
index ed4a5b0d..174f03d6 100644
--- a/pom.xml
+++ b/pom.xml
@@ -58,7 +58,7 @@
         <dependency>
             <groupId>io.getunleash</groupId>
             <artifactId>yggdrasil-engine</artifactId>
-            <version>0.1.2</version>
+            <version>0.2.2</version>
         </dependency>
         <dependency>
             <groupId>com.google.code.gson</groupId>
diff --git a/src/main/java/io/getunleash/DefaultUnleash.java b/src/main/java/io/getunleash/DefaultUnleash.java
index 3f68104d..dcccbb2e 100644
--- a/src/main/java/io/getunleash/DefaultUnleash.java
+++ b/src/main/java/io/getunleash/DefaultUnleash.java
@@ -2,6 +2,8 @@ package io.getunleash;
 
 import static io.getunleash.variant.Variant.DISABLED_VARIANT;
 
+import io.getunleash.engine.VariantDef;
+import io.getunleash.engine.WasmResponse;
 import io.getunleash.event.EventDispatcher;
 import io.getunleash.event.IsEnabledImpressionEvent;
 import io.getunleash.event.ToggleEvaluated;
@@ -111,14 +113,15 @@ public class DefaultUnleash implements Unleash {
 
         UnleashContext enhancedContext = context.applyStaticFields(config);
 
-        Boolean enabled = this.featureRepository.isEnabled(toggleName, enhancedContext);
+        WasmResponse<Boolean> response =
+                this.featureRepository.isEnabled(toggleName, enhancedContext);
+        Boolean enabled = response.value;
         if (enabled == null) {
             enabled = fallbackAction.test(toggleName, enhancedContext);
         }
 
-        this.metricService.countToggle(toggleName, enabled);
         eventDispatcher.dispatch(new ToggleEvaluated(toggleName, enabled));
-        if (this.featureRepository.shouldEmitImpressionEvent(toggleName)) {
+        if (response.impressionData) {
             eventDispatcher.dispatch(new IsEnabledImpressionEvent(toggleName, enabled, context));
         }
         return enabled;
@@ -142,15 +145,13 @@ public class DefaultUnleash implements Unleash {
     @Override
     public Variant getVariant(String toggleName, UnleashContext context, Variant defaultValue) {
         UnleashContext enhancedContext = context.applyStaticFields(config);
+        WasmResponse<VariantDef> response =
+                this.featureRepository.getVariant(toggleName, enhancedContext);
+        Optional<VariantDef> variantDef = Optional.ofNullable(response.value);
 
-        Variant variant =
-                YggdrasilAdapters.adapt(
-                        this.featureRepository.getVariant(toggleName, enhancedContext),
-                        defaultValue);
-        this.metricService.countToggle(toggleName, variant.isFeatureEnabled());
-        this.metricService.countVariant(toggleName, variant.getName());
+        Variant variant = YggdrasilAdapters.adapt(variantDef, defaultValue);
         eventDispatcher.dispatch(new ToggleEvaluated(toggleName, variant.isFeatureEnabled()));
-        if (this.featureRepository.shouldEmitImpressionEvent(toggleName)) {
+        if (response.impressionData) {
             eventDispatcher.dispatch(
                     new VariantImpressionEvent(
                             toggleName, variant.isFeatureEnabled(), context, variant.getName()));
@@ -199,11 +200,11 @@ public class DefaultUnleash implements Unleash {
                     .map(
                             toggleName -> {
                                 UnleashContext enhancedContext = context.applyStaticFields(config);
+                                WasmResponse<VariantDef> response =
+                                        featureRepository.getVariant(toggleName, enhancedContext);
+                                Optional<VariantDef> variantDef = Optional.of(response.value);
                                 Variant variant =
-                                        YggdrasilAdapters.adapt(
-                                                featureRepository.getVariant(
-                                                        toggleName, enhancedContext),
-                                                DISABLED_VARIANT);
+                                        YggdrasilAdapters.adapt(variantDef, DISABLED_VARIANT);
                                 return new EvaluatedToggle(
                                         toggleName, variant.isFeatureEnabled(), variant);
                             })
diff --git a/src/main/java/io/getunleash/EngineProxyImpl.java b/src/main/java/io/getunleash/EngineProxyImpl.java
index 4d4a7165..89a2835f 100644
--- a/src/main/java/io/getunleash/EngineProxyImpl.java
+++ b/src/main/java/io/getunleash/EngineProxyImpl.java
@@ -2,6 +2,7 @@ package io.getunleash;
 
 import io.getunleash.engine.UnleashEngine;
 import io.getunleash.engine.VariantDef;
+import io.getunleash.engine.WasmResponse;
 import io.getunleash.lang.Nullable;
 import io.getunleash.metric.UnleashMetricService;
 import io.getunleash.metric.UnleashMetricServiceImpl;
@@ -43,12 +44,12 @@ public class EngineProxyImpl implements EngineProxy {
     }
 
     @Override
-    public Boolean isEnabled(String toggleName, UnleashContext context) {
+    public WasmResponse<Boolean> isEnabled(String toggleName, UnleashContext context) {
         return this.featureRepository.isEnabled(toggleName, context);
     }
 
     @Override
-    public Optional<VariantDef> getVariant(String toggleName, UnleashContext context) {
+    public WasmResponse<VariantDef> getVariant(String toggleName, UnleashContext context) {
         return this.featureRepository.getVariant(toggleName, context);
     }
 
@@ -57,26 +58,11 @@ public class EngineProxyImpl implements EngineProxy {
         this.metricService.register(strategies);
     }
 
-    @Override
-    public void countToggle(String name, boolean enabled) {
-        this.metricService.countToggle(name, enabled);
-    }
-
-    @Override
-    public void countVariant(String name, String variantName) {
-        this.metricService.countVariant(name, variantName);
-    }
-
     @Override
     public Stream<FeatureDefinition> listKnownToggles() {
         return this.featureRepository.listKnownToggles();
     }
 
-    @Override
-    public boolean shouldEmitImpressionEvent(String toggleName) {
-        return this.featureRepository.shouldEmitImpressionEvent(toggleName);
-    }
-
     private static Map<String, Strategy> buildStrategyMap(@Nullable Strategy[] strategies) {
         Map<String, Strategy> map = new HashMap<>();
 
diff --git a/src/main/java/io/getunleash/FakeUnleash.java b/src/main/java/io/getunleash/FakeUnleash.java
index 2b57a085..5065ec6f 100644
--- a/src/main/java/io/getunleash/FakeUnleash.java
+++ b/src/main/java/io/getunleash/FakeUnleash.java
@@ -13,6 +13,7 @@ import java.util.stream.Stream;
 public class FakeUnleash implements Unleash {
     private boolean enableAll = false;
     private boolean disableAll = false;
+
     /**
      * @implNote This uses {@link Queue} instead of {@link List}, as there are concurrent queues,
      *     but no concurrent lists, in the jdk. This will never be drained. Only iterated over.
diff --git a/src/main/java/io/getunleash/metric/UnleashMetricService.java b/src/main/java/io/getunleash/metric/UnleashMetricService.java
index b701490b..d3602b7d 100644
--- a/src/main/java/io/getunleash/metric/UnleashMetricService.java
+++ b/src/main/java/io/getunleash/metric/UnleashMetricService.java
@@ -4,8 +4,4 @@ import java.util.Set;
 
 public interface UnleashMetricService {
     void register(Set<String> strategies);
-
-    void countToggle(String name, boolean enabled);
-
-    void countVariant(String name, String variantName);
 }
diff --git a/src/main/java/io/getunleash/metric/UnleashMetricServiceImpl.java b/src/main/java/io/getunleash/metric/UnleashMetricServiceImpl.java
index b50a9f52..eb8a7e16 100644
--- a/src/main/java/io/getunleash/metric/UnleashMetricServiceImpl.java
+++ b/src/main/java/io/getunleash/metric/UnleashMetricServiceImpl.java
@@ -2,7 +2,6 @@ package io.getunleash.metric;
 
 import io.getunleash.engine.MetricsBucket;
 import io.getunleash.engine.UnleashEngine;
-import io.getunleash.engine.YggdrasilError;
 import io.getunleash.util.Throttler;
 import io.getunleash.util.UnleashConfig;
 import io.getunleash.util.UnleashScheduledExecutor;
@@ -61,22 +60,17 @@ public class UnleashMetricServiceImpl implements UnleashMetricService {
     private Runnable sendMetrics() {
         return () -> {
             if (throttler.performAction()) {
-                try {
-                    MetricsBucket bucket = this.engine.getMetrics();
+                MetricsBucket bucket = this.engine.getMetrics();
 
-                    ClientMetrics metrics = new ClientMetrics(unleashConfig, bucket);
-                    int statusCode = metricSender.sendMetrics(metrics);
-                    if (statusCode >= 200 && statusCode < 400) {
-                        throttler.decrementFailureCountAndResetSkips();
-                    }
-                    if (statusCode >= 400) {
-                        throttler.handleHttpErrorCodes(statusCode);
-                    }
-                } catch (YggdrasilError e) {
-                    LOGGER.error(
-                            "Failed to retrieve metrics from the engine, this is a serious error, please report it",
-                            e);
+                ClientMetrics metrics = new ClientMetrics(unleashConfig, bucket);
+                int statusCode = metricSender.sendMetrics(metrics);
+                if (statusCode >= 200 && statusCode < 400) {
+                    throttler.decrementFailureCountAndResetSkips();
                 }
+                if (statusCode >= 400) {
+                    throttler.handleHttpErrorCodes(statusCode);
+                }
+
             } else {
                 throttler.skipped();
             }
@@ -90,22 +84,4 @@ public class UnleashMetricServiceImpl implements UnleashMetricService {
     protected int getFailures() {
         return this.throttler.getFailures();
     }
-
-    @Override
-    public void countToggle(String name, boolean enabled) {
-        try {
-            this.engine.countToggle(name, enabled);
-        } catch (YggdrasilError e) {
-            LOGGER.error("Failed to count toggle", e);
-        }
-    }
-
-    @Override
-    public void countVariant(String name, String variantName) {
-        try {
-            this.engine.countVariant(name, variantName);
-        } catch (YggdrasilError e) {
-            LOGGER.error("Failed to count variant", e);
-        }
-    }
 }
diff --git a/src/main/java/io/getunleash/repository/FeatureRepository.java b/src/main/java/io/getunleash/repository/FeatureRepository.java
index b2da2f8e..d988a89b 100644
--- a/src/main/java/io/getunleash/repository/FeatureRepository.java
+++ b/src/main/java/io/getunleash/repository/FeatureRepository.java
@@ -3,16 +3,14 @@ package io.getunleash.repository;
 import io.getunleash.FeatureDefinition;
 import io.getunleash.UnleashContext;
 import io.getunleash.engine.VariantDef;
-import java.util.Optional;
+import io.getunleash.engine.WasmResponse;
 import java.util.stream.Stream;
 
 public interface FeatureRepository {
 
-    Boolean isEnabled(String toggleName, UnleashContext context);
+    WasmResponse<Boolean> isEnabled(String toggleName, UnleashContext context);
 
-    Optional<VariantDef> getVariant(String toggleName, UnleashContext context);
+    WasmResponse<VariantDef> getVariant(String toggleName, UnleashContext context);
 
     Stream<FeatureDefinition> listKnownToggles();
-
-    boolean shouldEmitImpressionEvent(String toggleName);
 }
diff --git a/src/main/java/io/getunleash/repository/FeatureRepositoryImpl.java b/src/main/java/io/getunleash/repository/FeatureRepositoryImpl.java
index f5b13536..9e24ce8f 100644
--- a/src/main/java/io/getunleash/repository/FeatureRepositoryImpl.java
+++ b/src/main/java/io/getunleash/repository/FeatureRepositoryImpl.java
@@ -5,7 +5,7 @@ import io.getunleash.UnleashContext;
 import io.getunleash.UnleashException;
 import io.getunleash.engine.UnleashEngine;
 import io.getunleash.engine.VariantDef;
-import io.getunleash.engine.YggdrasilError;
+import io.getunleash.engine.WasmResponse;
 import io.getunleash.engine.YggdrasilInvalidInputException;
 import io.getunleash.event.ClientFeaturesResponse;
 import io.getunleash.event.EventDispatcher;
@@ -101,7 +101,7 @@ public class FeatureRepositoryImpl implements FeatureRepository {
         if (features.isPresent()) {
             try {
                 this.engine.takeState(features.get());
-            } catch (YggdrasilInvalidInputException | YggdrasilError e) {
+            } catch (YggdrasilInvalidInputException e) {
                 LOGGER.error("Error when initializing feature toggles", e);
                 eventDispatcher.dispatch(new UnleashException("Failed to read backup file:", e));
             }
@@ -160,7 +160,7 @@ public class FeatureRepositoryImpl implements FeatureRepository {
                     }
                 } catch (UnleashException e) {
                     handler.accept(e);
-                } catch (YggdrasilInvalidInputException | YggdrasilError e) {
+                } catch (YggdrasilInvalidInputException e) {
                     handler.accept(new UnleashException("Error when fetching features", e));
                 }
             } else {
@@ -178,43 +178,27 @@ public class FeatureRepositoryImpl implements FeatureRepository {
     }
 
     @Override
-    public Boolean isEnabled(String toggleName, UnleashContext context) {
+    public WasmResponse<Boolean> isEnabled(String toggleName, UnleashContext context) {
         try {
             return this.engine.isEnabled(toggleName, YggdrasilAdapters.adapt(context));
-        } catch (YggdrasilInvalidInputException | YggdrasilError e) {
+        } catch (YggdrasilInvalidInputException e) {
             LOGGER.error("Error when checking feature toggle {}", toggleName, e);
             return null;
         }
     }
 
     @Override
-    public Optional<VariantDef> getVariant(String toggleName, UnleashContext context) {
+    public WasmResponse<VariantDef> getVariant(String toggleName, UnleashContext context) {
         try {
-            return Optional.ofNullable(
-                    this.engine.getVariant(toggleName, YggdrasilAdapters.adapt(context)));
-        } catch (YggdrasilInvalidInputException | YggdrasilError e) {
+            return this.engine.getVariant(toggleName, YggdrasilAdapters.adapt(context));
+        } catch (YggdrasilInvalidInputException e) {
             LOGGER.error("Error when checking feature toggle {}", toggleName, e);
-            return Optional.empty();
+            return null;
         }
     }
 
     @Override
     public Stream<FeatureDefinition> listKnownToggles() {
-        try {
-            return this.engine.listKnownToggles().stream().map(FeatureDefinition::new);
-        } catch (YggdrasilError e) {
-            LOGGER.error("Error getting feature toggle definitions", e);
-            return Stream.empty();
-        }
-    }
-
-    @Override
-    public boolean shouldEmitImpressionEvent(String toggleName) {
-        try {
-            return this.engine.shouldEmitImpressionEvent(toggleName);
-        } catch (YggdrasilError e) {
-            LOGGER.error("Error checking impression event status on toggle{}", toggleName, e);
-            return false;
-        }
+        return this.engine.listKnownToggles().stream().map(FeatureDefinition::new);
     }
 }


Test patch:
diff --git a/src/test/java/io/getunleash/UnleashTest.java b/src/test/java/io/getunleash/UnleashTest.java
index 63817508..82087af5 100644
--- a/src/test/java/io/getunleash/UnleashTest.java
+++ b/src/test/java/io/getunleash/UnleashTest.java
@@ -4,6 +4,8 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
+import io.getunleash.engine.VariantDef;
+import io.getunleash.engine.WasmResponse;
 import io.getunleash.repository.ToggleBootstrapProvider;
 import io.getunleash.strategy.Strategy;
 import io.getunleash.util.UnleashConfig;
@@ -41,7 +43,8 @@ public class UnleashTest {
     @Test
     public void unknown_feature_should_use_default_setting() {
         EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class))).thenReturn(null);
+        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(false, null));
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
         boolean featureEnabled = unleash.isEnabled("test", true);
@@ -52,7 +55,8 @@ public class UnleashTest {
     @Test
     public void fallback_function_should_be_invoked_and_return_true() {
         EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class))).thenReturn(null);
+        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(false, null));
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
         BiPredicate<String, UnleashContext> fallbackAction = mock(BiPredicate.class);
@@ -65,7 +69,8 @@ public class UnleashTest {
     @Test
     public void fallback_function_should_be_invoked_also_with_context() {
         EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class))).thenReturn(null);
+        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(false, null));
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
         BiPredicate<String, UnleashContext> fallbackAction = mock(BiPredicate.class);
@@ -80,7 +85,8 @@ public class UnleashTest {
     @Test
     void fallback_function_should_be_invoked_and_return_false() {
         EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class))).thenReturn(null);
+        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(false, null));
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
         BiPredicate<String, UnleashContext> fallbackAction = mock(BiPredicate.class);
@@ -93,7 +99,8 @@ public class UnleashTest {
     @Test
     void fallback_function_should_not_be_called_when_toggle_is_defined() {
         EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class))).thenReturn(true);
+        when(engineProxy.isEnabled(eq("test"), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(true, true));
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
         BiPredicate<String, UnleashContext> fallbackAction = mock(BiPredicate.class);
@@ -142,7 +149,7 @@ public class UnleashTest {
         when(engineProxy.isEnabled(
                         eq("test"),
                         argThat(UnleashContext -> "111".equals(context.getUserId().orElse(null)))))
-                .thenReturn(Boolean.TRUE);
+                .thenReturn(new WasmResponse<Boolean>(true, true));
 
         UnleashConfig config =
                 createConfigBuilder().unleashContextProvider(contextProvider).build();
@@ -160,7 +167,7 @@ public class UnleashTest {
         when(engineProxy.isEnabled(
                         eq("test"),
                         argThat(UnleashContext -> "13".equals(context.getUserId().orElse(null)))))
-                .thenReturn(Boolean.TRUE);
+                .thenReturn(new WasmResponse<Boolean>(true, true));
 
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
@@ -172,7 +179,8 @@ public class UnleashTest {
         UnleashContext context = UnleashContext.builder().userId("13").build();
 
         EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.isEnabled(any(String.class), any(UnleashContext.class))).thenReturn(null);
+        when(engineProxy.isEnabled(any(String.class), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(false, null));
 
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
 
@@ -186,7 +194,7 @@ public class UnleashTest {
         Variant variant = new Variant("Chuck", "Norris", true);
         EngineProxy engineProxy = mock(EngineProxy.class);
         when(engineProxy.getVariant(any(String.class), any(UnleashContext.class)))
-                .thenReturn(Optional.empty());
+                .thenReturn(new WasmResponse<VariantDef>(true, null));
 
         Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
         Variant result = unleash.getVariant("test", context, variant);
@@ -194,20 +202,4 @@ public class UnleashTest {
         assertThat(result).isNotNull();
         assertThat(result.getName()).isEqualTo("Chuck");
     }
-
-    @Test
-    public void getting_variant_when_disabled_should_increment_no_counter() {
-
-        UnleashContext context = UnleashContext.builder().userId("1").build();
-
-        Variant variant = new Variant("Chuck", "Norris", true);
-        EngineProxy engineProxy = mock(EngineProxy.class);
-        when(engineProxy.getVariant(any(String.class), any(UnleashContext.class)))
-                .thenReturn(Optional.empty());
-
-        Unleash unleash = new DefaultUnleash(baseConfig, engineProxy);
-        Variant result = unleash.getVariant("test", context, variant);
-
-        verify(engineProxy, times(1)).countToggle("test", false);
-    }
 }
diff --git a/src/test/java/io/getunleash/event/ImpressionDataSubscriberTest.java b/src/test/java/io/getunleash/event/ImpressionDataSubscriberTest.java
index 8063254e..ff49ff7c 100644
--- a/src/test/java/io/getunleash/event/ImpressionDataSubscriberTest.java
+++ b/src/test/java/io/getunleash/event/ImpressionDataSubscriberTest.java
@@ -10,8 +10,8 @@ import io.getunleash.SynchronousTestExecutor;
 import io.getunleash.Unleash;
 import io.getunleash.UnleashContext;
 import io.getunleash.engine.VariantDef;
+import io.getunleash.engine.WasmResponse;
 import io.getunleash.util.UnleashConfig;
-import java.util.Optional;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
@@ -38,8 +38,8 @@ public class ImpressionDataSubscriberTest {
     public void noEventsIfImpressionDataIsNotEnabled() {
         String featureWithoutImpressionDataEnabled = "feature.with.no.impressionData";
         EngineProxy repo = Mockito.mock(EngineProxy.class);
-        when(repo.isEnabled(any(String.class), any(UnleashContext.class))).thenReturn(true);
-        when(repo.shouldEmitImpressionEvent(featureWithoutImpressionDataEnabled)).thenReturn(false);
+        when(repo.isEnabled(any(String.class), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(false, true));
         Unleash unleash = new DefaultUnleash(unleashConfig, repo);
 
         unleash.isEnabled(featureWithoutImpressionDataEnabled);
@@ -51,8 +51,8 @@ public class ImpressionDataSubscriberTest {
     public void isEnabledEventWhenImpressionDataIsEnabled() {
         String featureWithImpressionData = "feature.with.impressionData";
         EngineProxy repo = Mockito.mock(EngineProxy.class);
-        when(repo.isEnabled(any(String.class), any(UnleashContext.class))).thenReturn(true);
-        when(repo.shouldEmitImpressionEvent(featureWithImpressionData)).thenReturn(true);
+        when(repo.isEnabled(any(String.class), any(UnleashContext.class)))
+                .thenReturn(new WasmResponse<Boolean>(true, true));
         Unleash unleash = new DefaultUnleash(unleashConfig, repo);
 
         unleash.isEnabled(featureWithImpressionData);
@@ -65,9 +65,8 @@ public class ImpressionDataSubscriberTest {
         VariantDef mockVariant = Mockito.mock(VariantDef.class);
         String featureWithImpressionData = "feature.with.impressionData";
         EngineProxy repo = Mockito.mock(EngineProxy.class);
-        when(repo.shouldEmitImpressionEvent(featureWithImpressionData)).thenReturn(true);
         when(repo.getVariant(any(String.class), any(UnleashContext.class)))
-                .thenReturn(Optional.of(mockVariant));
+                .thenReturn(new WasmResponse<VariantDef>(true, mockVariant));
         Unleash unleash = new DefaultUnleash(unleashConfig, repo);
 
         unleash.getVariant(featureWithImpressionData);
diff --git a/src/test/java/io/getunleash/metric/UnleashMetricServiceImplTest.java b/src/test/java/io/getunleash/metric/UnleashMetricServiceImplTest.java
index 7f89da26..119cafb4 100644
--- a/src/test/java/io/getunleash/metric/UnleashMetricServiceImplTest.java
+++ b/src/test/java/io/getunleash/metric/UnleashMetricServiceImplTest.java
@@ -3,9 +3,11 @@ package io.getunleash.metric;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.*;
 
+import io.getunleash.engine.Context;
 import io.getunleash.engine.MetricsBucket;
 import io.getunleash.engine.UnleashEngine;
 import io.getunleash.engine.YggdrasilError;
+import io.getunleash.engine.YggdrasilInvalidInputException;
 import io.getunleash.util.UnleashConfig;
 import io.getunleash.util.UnleashScheduledExecutor;
 import java.time.LocalDateTime;
@@ -115,7 +117,8 @@ public class UnleashMetricServiceImplTest {
     }
 
     @Test
-    public void should_record_and_send_metrics() throws YggdrasilError {
+    public void should_record_and_send_metrics()
+            throws YggdrasilError, YggdrasilInvalidInputException {
         UnleashConfig config =
                 UnleashConfig.builder()
                         .appName("test")
@@ -130,54 +133,9 @@ public class UnleashMetricServiceImplTest {
 
         UnleashMetricService unleashMetricService =
                 new UnleashMetricServiceImpl(config, sender, executor, engine);
-        engine.countToggle("someToggle", true);
-        engine.countToggle("someToggle", false);
-        engine.countToggle("someToggle", true);
-        engine.countToggle("otherToggle", true);
 
-        // Call the sendMetricsCallback
-        ArgumentCaptor<Runnable> sendMetricsCallback = ArgumentCaptor.forClass(Runnable.class);
-        verify(executor).setInterval(sendMetricsCallback.capture(), anyLong(), anyLong());
-        sendMetricsCallback.getValue().run();
-
-        ArgumentCaptor<ClientMetrics> clientMetricsArgumentCaptor =
-                ArgumentCaptor.forClass(ClientMetrics.class);
-        verify(sender).sendMetrics(clientMetricsArgumentCaptor.capture());
-
-        ClientMetrics clientMetrics = clientMetricsArgumentCaptor.getValue();
-        MetricsBucket bucket = clientMetricsArgumentCaptor.getValue().getBucket();
-
-        assertThat(clientMetrics.getAppName()).isEqualTo(config.getAppName());
-        assertThat(clientMetrics.getEnvironment()).isEqualTo(config.getEnvironment());
-        assertThat(clientMetrics.getInstanceId()).isEqualTo(config.getInstanceId());
-        assertThat(clientMetrics.getConnectionId()).isEqualTo(config.getConnectionId());
-        assertThat(bucket.getStart()).isNotNull();
-        assertThat(bucket.getStop()).isNotNull();
-        assertThat(bucket.getToggles()).hasSize(2);
-        assertThat(bucket.getToggles().get("someToggle").getYes()).isEqualTo(2l);
-        assertThat(bucket.getToggles().get("someToggle").getNo()).isEqualTo(1l);
-    }
-
-    @Test
-    public void should_record_and_send_variant_metrics() throws YggdrasilError {
-        UnleashConfig config =
-                UnleashConfig.builder()
-                        .appName("test")
-                        .sendMetricsInterval(10)
-                        .unleashAPI("http://unleash.com")
-                        .build();
-
-        UnleashScheduledExecutor executor = mock(UnleashScheduledExecutor.class);
-        DefaultHttpMetricsSender sender = mock(DefaultHttpMetricsSender.class);
-        UnleashEngine engine = new UnleashEngine();
-
-        UnleashMetricService unleashMetricService =
-                new UnleashMetricServiceImpl(config, sender, executor, engine);
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v2");
-        engine.countVariant("someToggle", "disabled");
+        engine.isEnabled("someToggle", new Context());
+        engine.isEnabled("someToggle", new Context());
 
         // Call the sendMetricsCallback
         ArgumentCaptor<Runnable> sendMetricsCallback = ArgumentCaptor.forClass(Runnable.class);
@@ -192,19 +150,14 @@ public class UnleashMetricServiceImplTest {
         MetricsBucket bucket = clientMetricsArgumentCaptor.getValue().getBucket();
 
         assertThat(clientMetrics.getAppName()).isEqualTo(config.getAppName());
+        assertThat(clientMetrics.getEnvironment()).isEqualTo(config.getEnvironment());
         assertThat(clientMetrics.getInstanceId()).isEqualTo(config.getInstanceId());
         assertThat(clientMetrics.getConnectionId()).isEqualTo(config.getConnectionId());
         assertThat(bucket.getStart()).isNotNull();
         assertThat(bucket.getStop()).isNotNull();
         assertThat(bucket.getToggles()).hasSize(1);
-        assertThat(bucket.getToggles().get("someToggle").getVariants().get("v1").longValue())
-                .isEqualTo(3l);
-        assertThat(bucket.getToggles().get("someToggle").getVariants().get("v2").longValue())
-                .isEqualTo(1l);
-        assertThat(bucket.getToggles().get("someToggle").getVariants().get("disabled").longValue())
-                .isEqualTo(1l);
         assertThat(bucket.getToggles().get("someToggle").getYes()).isEqualTo(0l);
-        assertThat(bucket.getToggles().get("someToggle").getNo()).isEqualTo(0l);
+        assertThat(bucket.getToggles().get("someToggle").getNo()).isEqualTo(2l);
     }
 
     @Test
@@ -222,11 +175,6 @@ public class UnleashMetricServiceImplTest {
 
         UnleashMetricServiceImpl unleashMetricService =
                 new UnleashMetricServiceImpl(config, sender, executor, engine);
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v2");
-        engine.countVariant("someToggle", "disabled");
 
         // Call the sendMetricsCallback
         ArgumentCaptor<Runnable> sendMetricsCallback = ArgumentCaptor.forClass(Runnable.class);
@@ -295,11 +243,6 @@ public class UnleashMetricServiceImplTest {
 
         UnleashMetricServiceImpl unleashMetricService =
                 new UnleashMetricServiceImpl(config, sender, executor, engine);
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v2");
-        engine.countVariant("someToggle", "disabled");
 
         // Call the sendMetricsCallback
         ArgumentCaptor<Runnable> sendMetricsCallback = ArgumentCaptor.forClass(Runnable.class);
@@ -367,11 +310,6 @@ public class UnleashMetricServiceImplTest {
 
         UnleashMetricServiceImpl unleashMetricService =
                 new UnleashMetricServiceImpl(config, sender, executor, engine);
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v2");
-        engine.countVariant("someToggle", "disabled");
 
         // Call the sendMetricsCallback
         ArgumentCaptor<Runnable> sendMetricsCallback = ArgumentCaptor.forClass(Runnable.class);
@@ -405,11 +343,6 @@ public class UnleashMetricServiceImplTest {
 
         UnleashMetricServiceImpl unleashMetricService =
                 new UnleashMetricServiceImpl(config, sender, executor, engine);
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v1");
-        engine.countVariant("someToggle", "v2");
-        engine.countVariant("someToggle", "disabled");
 
         // Call the sendMetricsCallback
         ArgumentCaptor<Runnable> sendMetricsCallback = ArgumentCaptor.forClass(Runnable.class);