fix: paid stats test

supertokens · Sep 20, 2023 · ff0aad8 · ff0aad8
1 parent 05252f7
commit ff0aad8
Show file tree

Hide file tree

Showing 3 changed files with 90 additions and 2 deletions.
diff --git a/ee/src/main/java/io/supertokens/ee/EEFeatureFlag.java b/ee/src/main/java/io/supertokens/ee/EEFeatureFlag.java
@@ -452,7 +452,7 @@ private EE_FEATURES[] doServerCall(String licenseKey)
             if (Main.isTesting) {
                 licenseCheckRequests.add(json);
             }
-            ProcessState.getInstance(main).addState(ProcessState.PROCESS_STATE.LICENSE_KEY_CHECK_NETWORK_CALL, null);
+            ProcessState.getInstance(main).addState(ProcessState.PROCESS_STATE.LICENSE_KEY_CHECK_NETWORK_CALL, null, json);
             JsonObject licenseCheckResponse = HttpRequest.sendJsonPOSTRequest(this.main, REQUEST_ID,
                     "https://api.supertokens.io/0/st/license/check",
                     json, 10000, 10000, 0);

diff --git a/src/main/java/io/supertokens/ProcessState.java b/src/main/java/io/supertokens/ProcessState.java
@@ -16,6 +16,7 @@
 
 package io.supertokens;
 
+import com.google.gson.JsonObject;
 import io.supertokens.pluginInterface.multitenancy.TenantIdentifier;
 import io.supertokens.pluginInterface.multitenancy.exceptions.TenantOrAppNotFoundException;
 
@@ -51,8 +52,12 @@ public synchronized EventAndException getLastEventByName(PROCESS_STATE processSt
     }
 
     public synchronized void addState(PROCESS_STATE processState, Exception e) {
+        addState(processState, e, null);
+    }
+
+    public synchronized void addState(PROCESS_STATE processState, Exception e, JsonObject data) {
         if (Main.isTesting) {
-            history.add(new EventAndException(processState, e));
+            history.add(new EventAndException(processState, e, data));
         }
     }
 
@@ -97,11 +102,18 @@ public enum PROCESS_STATE {
 
     public static class EventAndException {
         public Exception exception;
+        public JsonObject data;
         PROCESS_STATE state;
 
         public EventAndException(PROCESS_STATE state, Exception e) {
             this.state = state;
             this.exception = e;
+            this.data = null;
+        }
+
+        public EventAndException(PROCESS_STATE state, Exception e, JsonObject data) {
+            this.state = state;
+            this.data = data;
         }
     }
 

diff --git a/src/test/java/io/supertokens/test/FeatureFlagTest.java b/src/test/java/io/supertokens/test/FeatureFlagTest.java
@@ -46,7 +46,9 @@
 import org.junit.*;
 import org.junit.rules.TestRule;
 
+import java.util.Arrays;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 import static org.junit.Assert.*;
@@ -744,4 +746,78 @@ public void testNetworkCallIsMadeInCoreInit() throws Exception {
         process.kill();
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
     }
+
+    private final String OPAQUE_KEY_WITH_DASHBOARD_FEATURE =
+            "EBy9Z4IRJ7BYyLP8AXxjq997o3RPaDekAE4CMGxduglUaEH9hugXzIduxvHIjpkFccVCZaHJIacMi8NJJg4I" +
+                    "=ruc3bZbT43QOLJbGu01cgACmVu2VOjQzFbT3lXiAKOR";
+
+    private final String OPAQUE_KEY_WITH_ACCOUNT_LINKING_FEATURE = "N2uEOdEzd1XZZ5VBSTGYaM7Ia4s8wAqRWFAxLqTYrB6GQ=" +
+            "vssOLo3c=PkFgcExkaXs=IA-d9UWccoNKsyUgNhOhcKtM1bjC5OLrYRpTAgN-2EbKYsQGGQRQHuUN4EO1V";
+
+    @Test
+    public void testPaidStatsContainsAllEnabledFeatures() throws Exception {
+        String[] args = {"../"};
+
+        EE_FEATURES[] allFeatures = EE_FEATURES.values();
+
+        List<EE_FEATURES> featuresToIgnore = List.of(EE_FEATURES.TEST);
+
+        String[] licenses = new String[]{
+                OPAQUE_KEY_WITH_MULTITENANCY_FEATURE,
+                OPAQUE_KEY_WITH_TOTP_FEATURE,
+                OPAQUE_KEY_WITH_DASHBOARD_FEATURE,
+                OPAQUE_KEY_WITH_ACCOUNT_LINKING_FEATURE
+        };
+
+        Set<EE_FEATURES> requiredFeatures = new HashSet<>();
+        requiredFeatures.addAll(Arrays.asList(allFeatures));
+        requiredFeatures.removeAll(featuresToIgnore);
+
+        Set<EE_FEATURES> foundFeatures = new HashSet<>();
+
+        for (String license : licenses) {
+            Utils.reset();
+            TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
+            process.startProcess();
+            assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
+
+            if (StorageLayer.getStorage(process.getProcess()).getType() != STORAGE_TYPE.SQL) {
+                return;
+            }
+
+            if (StorageLayer.isInMemDb(process.getProcess())) {
+                return;
+            }
+
+            // While adding license
+            TestMultitenancyAPIHelper.addLicense(license, process.getProcess());
+            assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.LICENSE_KEY_CHECK_NETWORK_CALL));
+            ProcessState.getInstance(process.getProcess()).clear();
+
+            process.kill(false);
+
+            // Restart core and check if the call was made during init
+            process = TestingProcessManager.start(args);
+            process.startProcess();
+            assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
+            ProcessState.EventAndException event = process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.LICENSE_KEY_CHECK_NETWORK_CALL);
+            assertNotNull(event);
+            assertNotNull(event.data);
+
+            JsonObject paidStatsObj = event.data.get("paidFeatureUsageStats").getAsJsonObject();
+            for (EE_FEATURES feature : allFeatures) {
+                if (featuresToIgnore.contains(feature)) {
+                    continue;
+                }
+                if (paidStatsObj.has(feature.toString())) {
+                    foundFeatures.add(feature);
+                }
+            }
+
+            process.kill();
+            assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
+        }
+
+        assertEquals(requiredFeatures, foundFeatures);
+    }
 }