fixed tests

amplitude · Jul 13, 2024 · 14e2004 · 14e2004
1 parent 85bcf02
commit 14e2004
Show file tree

Hide file tree

Showing 6 changed files with 563 additions and 396 deletions.
diff --git a/packages/node/src/local/cohort/fetcher.ts b/packages/node/src/local/cohort/fetcher.ts
@@ -49,15 +49,20 @@ export class CohortFetcher {
     this.logger = new ConsoleLogger(debug);
   }
 
+  static getKey(cohortId: string, lastModified?: number): string {
+    return `${cohortId}_${lastModified ? lastModified : ''}`;
+  }
+
   async fetch(
     cohortId: string,
     lastModified?: number,
   ): Promise<Cohort | undefined> {
     // This block may have async and awaits. No guarantee that executions are not interleaved.
     const unlock = await this.mutex.lock();
+    const key = CohortFetcher.getKey(cohortId, lastModified);
 
-    if (!this.inProgressCohorts[cohortId]) {
-      this.inProgressCohorts[cohortId] = this.executor.run(async () => {
+    if (!this.inProgressCohorts[key]) {
+      this.inProgressCohorts[key] = this.executor.run(async () => {
         this.logger.debug('Start downloading', cohortId);
         const cohort = await doWithBackoffFailLoudly<Cohort>(
           async () =>
@@ -73,22 +78,23 @@ export class CohortFetcher {
         )
           .then(async (cohort) => {
             const unlock = await this.mutex.lock();
-            delete this.inProgressCohorts[cohortId];
+            delete this.inProgressCohorts[key];
             unlock();
             return cohort;
           })
           .catch(async (err) => {
             const unlock = await this.mutex.lock();
-            delete this.inProgressCohorts[cohortId];
+            delete this.inProgressCohorts[key];
             unlock();
             throw err;
           });
-        this.logger.debug('Stop downloading', cohortId, cohort['cohortId']);
+        this.logger.debug('Stop downloading', cohortId);
         return cohort;
       });
     }
 
-    const cohortPromise = this.inProgressCohorts[cohortId];
+    const cohortPromise: Promise<Cohort | undefined> =
+      this.inProgressCohorts[key];
     unlock();
     return cohortPromise;
   }

diff --git a/packages/node/test/local/benchmark.test.ts b/packages/node/test/local/benchmark.test.ts
@@ -1,11 +1,30 @@
+import path from 'path';
+
+import * as dotenv from 'dotenv';
 import { Experiment } from 'src/factory';
 import { ExperimentUser } from 'src/types/user';
 
 import { measure } from './util/performance';
 
 const apiKey = 'server-Ed2doNl5YOblB5lRavQ9toj02arvHpMj';
 
-const client = Experiment.initializeLocal(apiKey, { debug: false });
+dotenv.config({ path: path.join(__dirname, '../../', '.env') });
+
+if (!process.env['API_KEY'] && !process.env['SECRET_KEY']) {
+  throw new Error(
+    'No env vars found. If running on local, have you created .env file correct environment variables? Checkout README.md',
+  );
+}
+
+const cohortConfig = {
+  apiKey: process.env['API_KEY'],
+  secretKey: process.env['SECRET_KEY'],
+};
+
+const client = Experiment.initializeLocal(apiKey, {
+  debug: false,
+  cohortConfig: cohortConfig,
+});
 
 beforeAll(async () => {
   await client.start();