Fix environment variable handling in jupyter paths unit tests (#16883)

* Checkpoint from VS Code for coding agent session

* Initial plan

* Update unit tests for enhanced kernel spec root paths

Co-authored-by: DonJayamanne <[email protected]>

* Add comprehensive unit tests for enhanced kernel spec root paths discovery

Co-authored-by: DonJayamanne <[email protected]>

* UpdatesÏ

* restore

* Update formatting

* Update cache key

* Updates

---------

Co-authored-by: Don Jayamanne <[email protected]>
Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: DonJayamanne <[email protected]>

Author: Copilot

src/kernels/raw/finder/jupyterPaths.node.ts

@@ -9,7 +9,7 @@ import * as uriPath from '../../../platform/vscode-path/resources';
 import { CancellationToken, Memento, Uri } from 'vscode';
 import { IFileSystem, IPlatformService } from '../../../platform/common/platform/types';
 import { IFileSystemNode } from '../../../platform/common/platform/types.node';
-import { ignoreLogging, logValue, logger } from '../../../platform/logging';
+import { logValue, logger } from '../../../platform/logging';
 import {
     IDisposableRegistry,
     IMemento,
@@ -39,15 +39,14 @@ const macJupyterRuntimePath = path.join('Library', 'Jupyter', 'runtime');
 
 export const baseKernelPath = path.join('share', 'jupyter', 'kernels');
 const CACHE_KEY_FOR_JUPYTER_KERNELSPEC_ROOT_PATH = 'CACHE_KEY_FOR_JUPYTER_KERNELSPEC_ROOT_PATH.';
-export const CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS = 'CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS_.';
+export const CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS = 'CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS_V2';
 
 /**
  * Finds locations to search for jupyter kernels.
  */
 @injectable()
 export class JupyterPaths {
     private cachedKernelSpecRootPath?: Promise<Uri | undefined>;
-    private cachedJupyterKernelPaths?: Promise<Uri[]>;
     private cachedJupyterPaths?: Promise<Uri[]>;
     private cachedDataDirs = new Map<string, Promise<Uri[]>>();
     constructor(
@@ -62,7 +61,6 @@ export class JupyterPaths {
     ) {
         this.envVarsProvider.onDidEnvironmentVariablesChange(
             () => {
-                this.cachedJupyterKernelPaths = undefined;
                 this.cachedJupyterPaths = undefined;
             },
             this,
@@ -330,24 +328,43 @@ export class JupyterPaths {
                 this.cachedKernelSpecRootPaths = undefined;
             }
         }, this);
+        promise
+            .then((paths) => {
+                if (this.cachedKernelSpecRootPaths?.promise === promise && paths.length) {
+                    this.updateCachedKernelSpecPaths(paths).catch(noop);
+                    this.cachedKernelSpecRootPaths.promise = Promise.resolve(paths);
+                }
+            })
+            .catch(noop);
         promise.finally(() => disposable.dispose()).catch(noop);
-        return promise;
+
+        // If we have something in cache use that.
+        const cached = this.getCachedKernelSpecPaths();
+        return cached.length ? cached : promise;
     }
     private async getKernelSpecRootPathsImpl(cancelToken: CancellationToken): Promise<Uri[]> {
+        const paths = new ResourceSet();
         // Paths specified in JUPYTER_PATH are supposed to come first in searching
-        const paths = new ResourceSet(await this.getJupyterPathKernelPaths(cancelToken));
+        // Use the complete data directory paths (equivalent to Python's jupyter_path("kernels"))
+        // This includes JUPYTER_PATH, user directories, environment directories, and system paths
+        const [dataDirs, kernelSpecRootpath] = await Promise.all([
+            this.getDataDirs({ resource: undefined }),
+            this.getKernelSpecRootPath()
+        ]);
         if (cancelToken.isCancellationRequested) {
             return [];
         }
-        if (this.platformService.isWindows) {
-            const winPath = await this.getKernelSpecRootPath();
-            if (cancelToken.isCancellationRequested) {
-                return [];
-            }
-            if (winPath) {
-                paths.add(winPath);
-            }
 
+        // Convert data directories to kernel spec directories by appending 'kernels' subdirectory
+        for (const dataDir of dataDirs) {
+            const kernelSpecDir = uriPath.joinPath(dataDir, 'kernels');
+            paths.add(kernelSpecDir);
+        }
+        if (kernelSpecRootpath) {
+            paths.add(kernelSpecRootpath);
+        }
+        // Add platform-specific additional paths that might not be covered by getDataDirs
+        if (this.platformService.isWindows) {
             if (process.env.PROGRAMDATA) {
                 paths.add(Uri.file(path.join(process.env.PROGRAMDATA, 'jupyter', 'kernels')));
             }
@@ -370,17 +387,6 @@ export class JupyterPaths {
         return Array.from(paths);
     }
 
-    private async getJupyterPathKernelPaths(@ignoreLogging() cancelToken?: CancellationToken): Promise<Uri[]> {
-        this.cachedJupyterKernelPaths =
-            this.cachedJupyterKernelPaths || this.getJupyterPathSubPaths(cancelToken, 'kernels');
-        this.cachedJupyterKernelPaths.then((value) => {
-            if (value.length > 0) {
-                this.updateCachedPaths(value).then(noop, noop);
-            }
-        }, noop);
-        return this.getCachedPaths().length > 0 ? this.getCachedPaths() : this.cachedJupyterKernelPaths;
-    }
-
     private async getJupyterPaths(cancelToken?: CancellationToken): Promise<Uri[]> {
         this.cachedJupyterPaths = this.cachedJupyterPaths || this.getJupyterPathSubPaths(cancelToken);
         return this.cachedJupyterPaths;
@@ -419,11 +425,11 @@ export class JupyterPaths {
         return Array.from(paths);
     }
 
-    private getCachedPaths(): Uri[] {
+    private getCachedKernelSpecPaths(): Uri[] {
         return this.globalState.get<string[]>(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, []).map((s) => Uri.parse(s));
     }
 
-    private async updateCachedPaths(paths: Uri[]) {
+    private async updateCachedKernelSpecPaths(paths: Uri[]) {
         const currentValue = this.globalState.get<string[]>(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, []);
         const newValue = paths.map(Uri.toString);
         if (currentValue.join(',') !== newValue.join(',')) {

src/kernels/raw/finder/jupyterPaths.node.unit.test.ts

@@ -6,7 +6,7 @@
 
 import * as sinon from 'sinon';
 import { assert } from 'chai';
-import { anything, deepEqual, instance, mock, when } from 'ts-mockito';
+import { anything, deepEqual, instance, mock, when, verify } from 'ts-mockito';
 import { CancellationTokenSource, ExtensionContext, Memento, Uri } from 'vscode';
 import { CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, JupyterPaths } from './jupyterPaths.node';
 import { dispose } from '../../../platform/common/utils/lifecycle';
@@ -48,6 +48,7 @@ suite('Jupyter Paths', () => {
     };
     const unixHomeDir = Uri.file('/users/username');
     const macHomeDir = Uri.file('/users/username');
+    const linuxJupyterPath = path.join('.local', 'share', 'jupyter', 'kernels');
     let cancelToken: CancellationTokenSource;
     suiteSetup(function () {
         if (isWeb()) {
@@ -322,8 +323,14 @@ suite('Jupyter Paths', () => {
         const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
         const winJupyterPath = path.join('AppData', 'Roaming', 'jupyter', 'kernels');
 
-        assert.strictEqual(paths.length, 1, `Expected 1 path, got ${paths.length}, ${JSON.stringify(paths)}`);
-        assert.strictEqual(paths[0].toString(), Uri.joinPath(windowsHomeDir, winJupyterPath).toString());
+        // New implementation returns data dirs + kernel spec root path
+        assert.strictEqual(paths.length, 2, `Expected 2 paths, got ${paths.length}, ${JSON.stringify(paths)}`);
+
+        // First path should be from data directory (.jupyter/data/kernels)
+        assert.strictEqual(paths[0].toString(), Uri.joinPath(windowsHomeDir, '.jupyter', 'data', 'kernels').toString());
+
+        // Second path should be the kernel spec root path
+        assert.strictEqual(paths[1].toString(), Uri.joinPath(windowsHomeDir, winJupyterPath).toString());
     });
 
     test('Get kernelspec root paths on Windows with JUPYTER_PATH env variable', async () => {
@@ -336,9 +343,17 @@ suite('Jupyter Paths', () => {
         const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
         const winJupyterPath = path.join('AppData', 'Roaming', 'jupyter', 'kernels');
 
-        assert.strictEqual(paths.length, 2, `Expected 2 path, got ${paths.length}, ${JSON.stringify(paths)}`);
+        // New implementation returns JUPYTER_PATH kernels + data dirs + kernel spec root path
+        assert.strictEqual(paths.length, 3, `Expected 3 paths, got ${paths.length}, ${JSON.stringify(paths)}`);
+
+        // First path should be from JUPYTER_PATH
         assert.strictEqual(paths[0].toString(), Uri.joinPath(Uri.file(__filename), 'kernels').toString());
-        assert.strictEqual(paths[1].toString(), Uri.joinPath(windowsHomeDir, winJupyterPath).toString());
+
+        // Second path should be from data directory (.jupyter/data/kernels)
+        assert.strictEqual(paths[1].toString(), Uri.joinPath(windowsHomeDir, '.jupyter', 'data', 'kernels').toString());
+
+        // Third path should be the kernel spec root path
+        assert.strictEqual(paths[2].toString(), Uri.joinPath(windowsHomeDir, winJupyterPath).toString());
     });
     test('Get kernelspec root paths on Windows with JUPYTER_PATH & ALLUSERSPROFILE env variable', async function () {
         when(platformService.osType).thenReturn(OSType.Windows);
@@ -351,12 +366,215 @@ suite('Jupyter Paths', () => {
         const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
         const winJupyterPath = path.join('AppData', 'Roaming', 'jupyter', 'kernels');
 
-        assert.strictEqual(paths.length, 3, `Expected 3 path, got ${paths.length}, ${JSON.stringify(paths)}`);
+        // New implementation returns JUPYTER_PATH kernels + data dirs + PROGRAMDATA + kernel spec root path
+        assert.strictEqual(paths.length, 4, `Expected 4 paths, got ${paths.length}, ${JSON.stringify(paths)}`);
+
+        // First path should be from JUPYTER_PATH
         assert.strictEqual(paths[0].toString(), Uri.joinPath(Uri.file(__filename), 'kernels').toString());
-        assert.strictEqual(paths[1].toString(), Uri.joinPath(windowsHomeDir, winJupyterPath).toString());
+
+        // Second path should be from data directory (.jupyter/data/kernels)
+        assert.strictEqual(paths[1].toString(), Uri.joinPath(windowsHomeDir, '.jupyter', 'data', 'kernels').toString());
+
+        // Third path should be from PROGRAMDATA
         assert.strictEqual(
             paths[2].toString(),
             Uri.file(path.join(allUserProfilePath, 'jupyter', 'kernels')).toString()
         );
+
+        // Fourth path should be the kernel spec root path
+        assert.strictEqual(paths[3].toString(), Uri.joinPath(windowsHomeDir, winJupyterPath).toString());
+    });
+
+    test('Get kernelspec root paths on Linux', async () => {
+        when(platformService.osType).thenReturn(OSType.Linux);
+        when(platformService.isWindows).thenReturn(false);
+        when(platformService.isMac).thenReturn(false);
+        when(platformService.homeDir).thenReturn(unixHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        // Clear environment variables that might affect the test
+        delete process.env['XDG_DATA_HOME'];
+        delete process.env['JUPYTER_DATA_DIR'];
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should include data dirs + system paths, and possibly kernel spec root path
+        assert.isAtLeast(paths.length, 3, `Expected at least 3 paths, got ${paths.length}, ${JSON.stringify(paths)}`);
+
+        // First path should be from data directory (without XDG_DATA_HOME, defaults to ~/.local/share/jupyter)
+        assert.strictEqual(
+            paths[0].toString(),
+            Uri.joinPath(unixHomeDir, '.local', 'share', 'jupyter', 'kernels').toString()
+        );
+
+        // Should include system paths
+        const pathStrings = paths.map((p) => p.toString());
+        assert.include(pathStrings, Uri.file('/usr/share/jupyter/kernels').toString());
+        assert.include(pathStrings, Uri.file('/usr/local/share/jupyter/kernels').toString());
+
+        // May include kernel spec root path if available
+        const hasKernelSpecRootPath = pathStrings.some((p) => p.includes(linuxJupyterPath));
+        if (hasKernelSpecRootPath) {
+            assert.include(pathStrings, Uri.joinPath(unixHomeDir, linuxJupyterPath).toString());
+        }
+    });
+
+    test('Get kernelspec root paths on macOS', async () => {
+        when(platformService.osType).thenReturn(OSType.OSX);
+        when(platformService.isWindows).thenReturn(false);
+        when(platformService.isMac).thenReturn(true);
+        when(platformService.homeDir).thenReturn(macHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        // Clear environment variables that might affect the test
+        delete process.env['XDG_DATA_HOME'];
+        delete process.env['JUPYTER_DATA_DIR'];
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should include at least 3 paths: data dir + system paths
+        assert.isAtLeast(paths.length, 3, `Expected at least 3 paths, got ${paths.length}, ${JSON.stringify(paths)}`);
+
+        // First path should be from data directory (macOS uses ~/Library/Jupyter)
+        assert.strictEqual(paths[0].toString(), Uri.joinPath(macHomeDir, 'Library', 'Jupyter', 'kernels').toString());
+
+        // Should include system paths
+        const pathStrings = paths.map((p) => p.toString());
+        assert.include(pathStrings, Uri.file('/usr/share/jupyter/kernels').toString());
+        assert.include(pathStrings, Uri.file('/usr/local/share/jupyter/kernels').toString());
+    });
+
+    test('Get kernelspec root paths on Linux with Python interpreter', async () => {
+        when(platformService.osType).thenReturn(OSType.Linux);
+        when(platformService.isWindows).thenReturn(false);
+        when(platformService.isMac).thenReturn(false);
+        when(platformService.homeDir).thenReturn(unixHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        // Mock getDataDirs to return additional interpreter-specific paths
+        const mockDataDirs = [
+            Uri.joinPath(unixHomeDir, '.local', 'share', 'jupyter'),
+            Uri.joinPath(Uri.file(sysPrefix), 'share', 'jupyter')
+        ];
+        sinon.stub(jupyterPaths, 'getDataDirs').resolves(mockDataDirs);
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should include interpreter-specific data dirs converted to kernel paths
+        const pathStrings = paths.map((p) => p.toString());
+        assert.include(pathStrings, Uri.joinPath(unixHomeDir, '.local', 'share', 'jupyter', 'kernels').toString());
+        assert.include(pathStrings, Uri.joinPath(Uri.file(sysPrefix), 'share', 'jupyter', 'kernels').toString());
+
+        sinon.restore();
+    });
+
+    test('Get kernelspec root paths handles cancellation token', async () => {
+        when(platformService.osType).thenReturn(OSType.Windows);
+        when(platformService.homeDir).thenReturn(windowsHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        // Cancel the token immediately
+        cancelToken.cancel();
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should return empty array when cancelled
+        assert.strictEqual(paths.length, 0, `Expected empty array when cancelled, got ${paths.length} paths`);
+    });
+
+    test('Get kernelspec root paths handles missing home directory gracefully', async () => {
+        when(platformService.osType).thenReturn(OSType.Linux);
+        when(platformService.isWindows).thenReturn(false);
+        when(platformService.isMac).thenReturn(false);
+        when(platformService.homeDir).thenReturn(undefined); // No home directory
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should still return system paths even without home directory
+        const pathStrings = paths.map((p) => p.toString());
+        assert.include(pathStrings, Uri.file('/usr/share/jupyter/kernels').toString());
+        assert.include(pathStrings, Uri.file('/usr/local/share/jupyter/kernels').toString());
+    });
+
+    test('Get kernelspec root paths deduplicates paths', async () => {
+        when(platformService.osType).thenReturn(OSType.Windows);
+        when(platformService.homeDir).thenReturn(windowsHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        // Create a scenario where paths might be duplicated
+        const duplicatePath = Uri.joinPath(windowsHomeDir, '.jupyter', 'data');
+        const mockDataDirs = [
+            duplicatePath,
+            duplicatePath, // Duplicate
+            Uri.joinPath(windowsHomeDir, 'AppData', 'Roaming', 'jupyter')
+        ];
+        sinon.stub(jupyterPaths, 'getDataDirs').resolves(mockDataDirs);
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should not contain duplicate paths
+        const pathStrings = paths.map((p) => p.toString());
+        const uniquePaths = [...new Set(pathStrings)];
+        assert.strictEqual(pathStrings.length, uniquePaths.length, 'Paths should be deduplicated');
+
+        sinon.restore();
+    });
+
+    test('Get kernelspec root paths with cached data', async () => {
+        const cachedPaths = [
+            Uri.joinPath(windowsHomeDir, 'cached1', 'kernels').toString(),
+            Uri.joinPath(windowsHomeDir, 'cached2', 'kernels').toString()
+        ];
+        when(platformService.osType).thenReturn(OSType.Windows);
+        when(platformService.homeDir).thenReturn(windowsHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn(cachedPaths);
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Should return cached data if available
+        assert.strictEqual(paths.length, 2);
+        assert.strictEqual(paths[0].toString(), cachedPaths[0]);
+        assert.strictEqual(paths[1].toString(), cachedPaths[1]);
+    });
+
+    test('Get kernelspec root paths with JUPYTER_PATH on Linux', async () => {
+        when(platformService.osType).thenReturn(OSType.Linux);
+        when(platformService.isWindows).thenReturn(false);
+        when(platformService.isMac).thenReturn(false);
+        when(platformService.homeDir).thenReturn(unixHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        const jupyter_Paths = ['/custom/jupyter/path1', '/custom/jupyter/path2'];
+        process.env['JUPYTER_PATH'] = jupyter_Paths.join(path.delimiter);
+
+        const paths = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // First paths should be from JUPYTER_PATH with 'kernels' appended
+        assert.isAtLeast(paths.length, 2, `Expected at least 2 paths, got ${paths.length}`);
+        assert.strictEqual(paths[0].toString(), Uri.joinPath(Uri.file(jupyter_Paths[0]), 'kernels').toString());
+        assert.strictEqual(paths[1].toString(), Uri.joinPath(Uri.file(jupyter_Paths[1]), 'kernels').toString());
+    });
+
+    test('Enhanced caching behavior works correctly', async () => {
+        when(platformService.osType).thenReturn(OSType.Windows);
+        when(platformService.homeDir).thenReturn(windowsHomeDir);
+        when(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).thenReturn([]);
+
+        // First call
+        const paths1 = await jupyterPaths.getKernelSpecRootPaths(cancelToken.token);
+
+        // Verify caching method is called
+        verify(memento.get(CACHE_KEY_FOR_JUPYTER_KERNEL_PATHS, anything())).atLeast(1);
+
+        // Second call should use cached result (simulate by using same token)
+        const newCancelToken = new CancellationTokenSource();
+        disposables.push(newCancelToken);
+        const paths2 = await jupyterPaths.getKernelSpecRootPaths(newCancelToken.token);
+
+        assert.strictEqual(paths1.length, paths2.length);
+        paths1.forEach((path, index) => {
+            assert.strictEqual(path.toString(), paths2[index].toString());
+        });
     });
 });