From a49302de9d7b3c39c97ae343e7a41c7a8716e9c4 Mon Sep 17 00:00:00 2001
From: Ben White <ben@posthog.com>
Date: Tue, 3 Dec 2024 17:48:24 +0100
Subject: [PATCH] Fix tests

---
 src/__tests__/web-experiments.test.ts | 31 +++++++++++++++++----------
 src/web-experiments.ts                |  2 +-
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/__tests__/web-experiments.test.ts b/src/__tests__/web-experiments.test.ts
index 2e5df4f11..3f504c30d 100644
--- a/src/__tests__/web-experiments.test.ts
+++ b/src/__tests__/web-experiments.test.ts
@@ -1,6 +1,6 @@
 import { WebExperiments } from '../web-experiments'
 import { PostHog } from '../posthog-core'
-import { DecideResponse, PostHogConfig } from '../types'
+import { PostHogConfig } from '../types'
 import { PostHogPersistence } from '../posthog-persistence'
 import { WebExperiment } from '../web-experiments-types'
 import { RequestRouter } from '../utils/request-router'
@@ -111,7 +111,10 @@ describe('Web Experimentation', () => {
         },
     } as unknown as WebExperiment
 
+    const simulateFeatureFlags: jest.Mock = jest.fn()
+
     beforeEach(() => {
+        let cachedFlags = {}
         persistence = { props: {}, register: jest.fn() } as unknown as PostHogPersistence
         posthog = makePostHog({
             config: {
@@ -129,6 +132,14 @@ describe('Web Experimentation', () => {
                 .mockImplementation(({ callback }) => callback({ statusCode: 200, json: experimentsResponse })),
             consent: { isOptedOut: () => true } as unknown as ConsentManager,
             onFeatureFlags: jest.fn(),
+            getFeatureFlag: (key: string) => {
+                return cachedFlags[key]
+            },
+        })
+
+        simulateFeatureFlags.mockImplementation((flags) => {
+            cachedFlags = flags
+            webExperiment.onFeatureFlags(Object.keys(flags))
         })
 
         posthog.requestRouter = new RequestRouter(posthog)
@@ -171,11 +182,10 @@ describe('Web Experimentation', () => {
     function assertElementChanged(variant: string, expectedProperty: string, value: string) {
         const elParent = createTestDocument()
         webExperiment = new WebExperiments(posthog)
-        webExperiment.onRemoteConfig({
-            featureFlags: {
-                'signup-button-test': variant,
-            },
-        } as unknown as DecideResponse)
+
+        simulateFeatureFlags({
+            'signup-button-test': variant,
+        })
 
         switch (expectedProperty) {
             case 'css':
@@ -201,11 +211,10 @@ describe('Web Experimentation', () => {
             webExperiment._is_bot = () => true
             const elParent = createTestDocument()
 
-            webExperiment.onRemoteConfig({
-                featureFlags: {
-                    'signup-button-test': 'Sign me up',
-                },
-            } as unknown as DecideResponse)
+            simulateFeatureFlags({
+                'signup-button-test': 'Sign me up',
+            })
+
             expect(elParent.innerText).toEqual('original')
         })
     })
diff --git a/src/web-experiments.ts b/src/web-experiments.ts
index 244b10bfe..74bbc904b 100644
--- a/src/web-experiments.ts
+++ b/src/web-experiments.ts
@@ -114,7 +114,7 @@ export class WebExperiments {
                         this._flagToExperiments?.set(webExperiment.feature_flag_key, webExperiment)
                     }
 
-                    const selectedVariant = this.instance.featureFlags.getFeatureFlag(webExperiment.feature_flag_key)
+                    const selectedVariant = this.instance.getFeatureFlag(webExperiment.feature_flag_key)
                     if (isString(selectedVariant) && webExperiment.variants[selectedVariant]) {
                         this.applyTransforms(
                             webExperiment.name,