more tests for bandits polling

Author: maya-the-mango

src/configuration-requestor.spec.ts

@@ -230,6 +230,77 @@ describe('ConfigurationRequestor', () => {
         expect(fetchSpy).toHaveBeenCalledTimes(3); // Once just for UFC, bandits should be skipped
       });
 
+      const warmStartBanditReference = {
+        modelVersion: 'warm start',
+        flagVariations: [
+          {
+            key: 'warm_start_bandit',
+            flagKey: 'warm_start_bandit_flag',
+            variationKey: 'warm_start_bandit',
+            variationValue: 'warm_start_bandit',
+          },
+        ],
+      };
+
+      const warmStartBanditParameters = {
+        banditKey: 'warm_start_bandit',
+        modelName: 'pigeon',
+        modelVersion: 'warm start',
+        modelData: {
+          gamma: 1.0,
+          defaultActionScore: 0.0,
+          actionProbabilityFloor: 0.0,
+          coefficients: {},
+        },
+      };
+
+      const coldStartBanditParameters = {
+        banditKey: 'cold_start_bandit',
+        modelName: 'falcon',
+        modelVersion: 'cold start',
+        modelData: {
+          gamma: 1.0,
+          defaultActionScore: 0.0,
+          actionProbabilityFloor: 0.0,
+          coefficients: {},
+        },
+      };
+
+      function expectBanditToBeInModelStore(
+        store: IConfigurationStore<BanditParameters>,
+        banditKey: string,
+        expectedBanditParameters: BanditParameters,
+      ) {
+        const bandit = store.get(banditKey);
+        expect(bandit).toBeTruthy();
+        expect(bandit?.banditKey).toBe(expectedBanditParameters.banditKey);
+        expect(bandit?.modelVersion).toBe(expectedBanditParameters.modelVersion);
+        expect(bandit?.modelName).toBe(expectedBanditParameters.modelName);
+        expect(bandit?.modelData.gamma).toBe(expectedBanditParameters.modelData.gamma);
+        expect(bandit?.modelData.defaultActionScore).toBe(
+          expectedBanditParameters.modelData.defaultActionScore,
+        );
+        expect(bandit?.modelData.actionProbabilityFloor).toBe(
+          expectedBanditParameters.modelData.actionProbabilityFloor,
+        );
+        expect(bandit?.modelData.coefficients).toStrictEqual(
+          expectedBanditParameters.modelData.coefficients,
+        );
+      }
+
+      function injectWarmStartBanditToResponseByUrl(
+        url: string,
+        response: IUniversalFlagConfigResponse | IBanditParametersResponse,
+      ) {
+        if (url.includes('config') && 'banditReferences' in response) {
+          response.banditReferences.warm_start_bandit = warmStartBanditReference;
+        }
+
+        if (url.includes('bandits') && 'bandits' in response) {
+          response.bandits.warm_start_bandit = warmStartBanditParameters;
+        }
+      }
+
       it('Should fetch bandits if new bandit references model versions appeared', async () => {
         let updateUFC = false;
         await configurationRequestor.fetchAndStoreConfigurations();
@@ -244,34 +315,7 @@ describe('ConfigurationRequestor', () => {
           const response = readMockUFCResponse(responseFile);
 
           if (updateUFC === true) {
-            // this if is needed to appease linter
-            if (url.includes('config') && 'banditReferences' in response) {
-              response.banditReferences.warm_start = {
-                modelVersion: 'warm start',
-                flagVariations: [
-                  {
-                    key: 'warm_start_bandit',
-                    flagKey: 'warm_start_bandit_flag',
-                    variationKey: 'warm_start_bandit',
-                    variationValue: 'warm_start_bandit',
-                  },
-                ],
-              };
-            }
-
-            if (url.includes('bandits') && 'bandits' in response) {
-              response.bandits.warm_start = {
-                banditKey: 'warm_start_bandit',
-                modelName: 'pigeon',
-                modelVersion: 'warm start',
-                modelData: {
-                  gamma: 1.0,
-                  defaultActionScore: 0.0,
-                  actionProbabilityFloor: 0.0,
-                  coefficients: {},
-                },
-              };
-            }
+            injectWarmStartBanditToResponseByUrl(url, response);
           }
           return response;
         };
@@ -282,15 +326,11 @@ describe('ConfigurationRequestor', () => {
         expect(fetchSpy).toHaveBeenCalledTimes(2); // 2 because fetchSpy was re-initiated, 1UFC and 1bandits
 
         // let's check if warm start was hydrated properly!
-        const warm_start_bandit = banditModelStore.get('warm_start');
-        expect(warm_start_bandit).toBeTruthy();
-        expect(warm_start_bandit?.banditKey).toBe('warm_start_bandit');
-        expect(warm_start_bandit?.modelVersion).toBe('warm start');
-        expect(warm_start_bandit?.modelName).toBe('pigeon');
-        expect(warm_start_bandit?.modelData.gamma).toBe(1);
-        expect(warm_start_bandit?.modelData.defaultActionScore).toBe(0);
-        expect(warm_start_bandit?.modelData.actionProbabilityFloor).toBe(0);
-        expect(warm_start_bandit?.modelData.coefficients).toStrictEqual({});
+        expectBanditToBeInModelStore(
+          banditModelStore,
+          'warm_start_bandit',
+          warmStartBanditParameters,
+        );
       });
 
       it('Should not fetch bandits if bandit references model versions shrunk', async () => {
@@ -316,15 +356,67 @@ describe('ConfigurationRequestor', () => {
         expect(fetchSpy).toHaveBeenCalledTimes(1); // only once for UFC
 
         // cold start should still be in memory
-        const warm_start_bandit = banditModelStore.get('cold_start_bandit');
-        expect(warm_start_bandit).toBeTruthy();
-        expect(warm_start_bandit?.banditKey).toBe('cold_start_bandit');
-        expect(warm_start_bandit?.modelVersion).toBe('cold start');
-        expect(warm_start_bandit?.modelName).toBe('falcon');
-        expect(warm_start_bandit?.modelData.gamma).toBe(1);
-        expect(warm_start_bandit?.modelData.defaultActionScore).toBe(0);
-        expect(warm_start_bandit?.modelData.actionProbabilityFloor).toBe(0);
-        expect(warm_start_bandit?.modelData.coefficients).toStrictEqual({});
+        expectBanditToBeInModelStore(
+          banditModelStore,
+          'cold_start_bandit',
+          coldStartBanditParameters,
+        );
+      });
+
+      /**
+       * 1. initial call - 1 fetch for ufc 1 for bandits
+       * 2. 2nd call - 1 fetch for ufc only; bandits unchanged
+       * 3. 3rd call - new bandit ref injected to UFC; 2 fetches, because new bandit appeared
+       * 4. 4th call - we remove a bandit from ufc; 1 fetch because there is no need to update.
+       *    The bandit removed from UFC should still be in memory.
+       **/
+      it('should fetch bandits based on banditReference change in UFC', async () => {
+        let injectWarmStart = false;
+        let removeColdStartBandit = false;
+        await configurationRequestor.fetchAndStoreConfigurations();
+        expect(fetchSpy).toHaveBeenCalledTimes(2);
+
+        await configurationRequestor.fetchAndStoreConfigurations();
+        expect(fetchSpy).toHaveBeenCalledTimes(3);
+
+        const customResponseMockGenerator = (url: string) => {
+          const responseFile = url.includes('bandits')
+            ? MOCK_BANDIT_MODELS_RESPONSE_FILE
+            : MOCK_FLAGS_WITH_BANDITS_RESPONSE_FILE;
+          const response = readMockUFCResponse(responseFile);
+          if (injectWarmStart === true) {
+            injectWarmStartBanditToResponseByUrl(url, response);
+          } else if (
+            removeColdStartBandit === true &&
+            'banditReferences' in response &&
+            url.includes('config')
+          ) {
+            delete response.banditReferences.cold_start_bandit;
+          }
+          return response;
+        };
+        injectWarmStart = true;
+        initiateFetchSpy(customResponseMockGenerator);
+
+        await configurationRequestor.fetchAndStoreConfigurations();
+        expect(fetchSpy).toHaveBeenCalledTimes(2);
+        expectBanditToBeInModelStore(
+          banditModelStore,
+          'warm_start_bandit',
+          warmStartBanditParameters,
+        );
+
+        injectWarmStart = false;
+        removeColdStartBandit = true;
+        initiateFetchSpy(customResponseMockGenerator);
+        await configurationRequestor.fetchAndStoreConfigurations();
+        expect(fetchSpy).toHaveBeenCalledTimes(1);
+
+        expectBanditToBeInModelStore(
+          banditModelStore,
+          'cold_start_bandit',
+          coldStartBanditParameters,
+        );
       });
     });
   });