Add tests for environments collection (#17965)

* Add test for environments collection

* Add test to handle updates from downstream locators

* Add more tests

* Fix tests

* Minor stuff

* Apply suggestions from code review

Co-authored-by: Kim-Adeline Miguel <[email protected]>

* Add test for missing branch

Co-authored-by: Kim-Adeline Miguel <[email protected]>

Author: Kartik Raj

src/client/pythonEnvironments/base/info/env.ts

@@ -3,6 +3,7 @@
 
 import { cloneDeep } from 'lodash';
 import * as path from 'path';
+import { Uri } from 'vscode';
 import { getArchitectureDisplayName } from '../../../common/platform/registry';
 import { normalizeFilename } from '../../../common/utils/filesystem';
 import { Architecture } from '../../../common/utils/platform';
@@ -29,6 +30,7 @@ export function buildEnvInfo(init?: {
     source?: PythonEnvSource[];
     display?: string;
     sysPrefix?: string;
+    searchLocation?: Uri;
 }): PythonEnvInfo {
     const env = {
         name: init?.name ?? '',
@@ -90,6 +92,7 @@ function updateEnv(
         executable?: string;
         location?: string;
         version?: PythonVersion;
+        searchLocation?: Uri;
     },
 ): void {
     if (updates.kind !== undefined) {
@@ -104,6 +107,9 @@ function updateEnv(
     if (updates.version !== undefined) {
         env.version = updates.version;
     }
+    if (updates.searchLocation !== undefined) {
+        env.searchLocation = updates.searchLocation;
+    }
 }
 
 /**

src/client/pythonEnvironments/base/locators/composite/envsCollectionCache.ts

@@ -52,10 +52,10 @@ export interface IEnvsCollectionCache {
     validateCache(): Promise<void>;
 }
 
-type PythonEnvCompleteInfo = { hasCompleteInfo?: boolean } & PythonEnvInfo;
+export type PythonEnvCompleteInfo = { hasCompleteInfo?: boolean } & PythonEnvInfo;
 
 interface IPersistentStorage {
-    load(): Promise<PythonEnvInfo[] | undefined>;
+    load(): Promise<PythonEnvInfo[]>;
     store(envs: PythonEnvInfo[]): Promise<void>;
 }
 
@@ -118,7 +118,7 @@ export class PythonEnvInfoCache extends PythonEnvsWatcher<PythonEnvCollectionCha
     }
 
     public async clearAndReloadFromStorage(): Promise<void> {
-        this.envs = (await this.persistentStorage.load()) ?? this.envs;
+        this.envs = await this.persistentStorage.load();
     }
 
     public async flush(): Promise<void> {

src/client/pythonEnvironments/base/locators/composite/envsCollectionService.ts

@@ -21,8 +21,8 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
     /** Keeps track of ongoing refreshes for various queries. */
     private refreshPromises = new Map<PythonLocatorQuery | undefined, Promise<void>>();
 
-    /** Keeps track of whether there are any scheduled refreshes other than the ongoing one for various queries. */
-    private scheduledRefreshes = new Map<PythonLocatorQuery | undefined, boolean>();
+    /** Keeps track of scheduled refreshes other than the ongoing one for various queries. */
+    private scheduledRefreshes = new Map<PythonLocatorQuery | undefined, Promise<void>>();
 
     private readonly refreshStarted = new EventEmitter<void>();
 
@@ -40,14 +40,14 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
         super();
         this.locator.onChanged((event) => {
             const query = undefined; // We can also form a query based on the event, but skip that for simplicity.
-            const isNewRefreshScheduled = this.scheduledRefreshes.get(query);
-            if (isNewRefreshScheduled) {
-                // If there is already a new refresh scheduled for the query, no need to start another one.
-                return;
+            let scheduledRefresh = this.scheduledRefreshes.get(query);
+            // If there is no refresh scheduled for the query, start a new one.
+            if (!scheduledRefresh) {
+                scheduledRefresh = this.scheduleNewRefresh(query);
             }
-            this.scheduleNewRefresh(query).then(() => {
+            scheduledRefresh.then(() => {
                 // Once refresh of cache is complete, notify changes.
-                this.fire({ type: event.type, searchLocation: event.searchLocation });
+                this.fire(event);
             });
         });
         this.cache.onChanged((e) => {
@@ -73,7 +73,8 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
     public getEnvs(query?: PythonLocatorQuery): PythonEnvInfo[] {
         const cachedEnvs = this.cache.getAllEnvs();
         if (cachedEnvs.length === 0 && this.refreshPromises.size === 0) {
-            traceError('Refresh should have already been triggered when activating discovery component');
+            // We expect a refresh to already be triggered when activating discovery component.
+            traceError('No python is installed or a refresh has not already been triggered');
             this.triggerRefresh().ignoreErrors();
         }
         return query ? cachedEnvs.filter(getQueryFilter(query)) : cachedEnvs;
@@ -159,13 +160,18 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
      * Ensure we trigger a fresh refresh for the query after the current refresh (if any) is done.
      */
     private async scheduleNewRefresh(query?: PythonLocatorQuery): Promise<void> {
-        this.scheduledRefreshes.set(query, true);
-        const refreshPromise = this.getRefreshPromiseForQuery(query) ?? Promise.resolve();
-        const nextRefreshPromise = refreshPromise.then(() => {
-            // No more scheduled refreshes for this query as we're about to start the scheduled one.
-            this.scheduledRefreshes.set(query, false);
-            this.startRefresh(query);
-        });
+        const refreshPromise = this.getRefreshPromiseForQuery(query);
+        let nextRefreshPromise: Promise<void>;
+        if (!refreshPromise) {
+            nextRefreshPromise = this.startRefresh(query);
+        } else {
+            nextRefreshPromise = refreshPromise.then(() => {
+                // No more scheduled refreshes for this query as we're about to start the scheduled one.
+                this.scheduledRefreshes.delete(query);
+                this.startRefresh(query);
+            });
+            this.scheduledRefreshes.set(query, nextRefreshPromise);
+        }
         return nextRefreshPromise;
     }
 }

src/client/pythonEnvironments/index.ts

@@ -161,7 +161,7 @@ function createWorkspaceLocator(ext: ExtensionState): WorkspaceLocators<BasicEnv
 }
 
 async function createCollectionCache(ext: ExtensionState): Promise<IEnvsCollectionCache> {
-    const storage = getGlobalStorage<PythonEnvInfo[]>(ext.context, 'PYTHON_ENV_INFO_CACHE');
+    const storage = getGlobalStorage<PythonEnvInfo[]>(ext.context, 'PYTHON_ENV_INFO_CACHE', []);
     const cache = await createCache({
         load: async () => storage.get(),
         store: async (e) => storage.set(e),

src/test/pythonEnvironments/base/common.ts

@@ -87,7 +87,7 @@ export class SimpleLocator<I = PythonEnvInfo> extends Locator<I> {
         public callbacks: {
             resolve?: null | ((env: PythonEnvInfo) => Promise<PythonEnvInfo | undefined>);
             before?: Promise<void>;
-            after?: Promise<void>;
+            after?(): Promise<void>;
             onUpdated?: Event<PythonEnvUpdatedEvent<I> | null>;
             beforeEach?(e: I): Promise<void>;
             afterEach?(e: I): Promise<void>;
@@ -137,7 +137,7 @@ export class SimpleLocator<I = PythonEnvInfo> extends Locator<I> {
                 }
             }
             if (callbacks?.after !== undefined) {
-                await callbacks.after;
+                await callbacks.after();
             }
             deferred.resolve();
         })();