venvUtils.getVenvForWorkspace.unit.test.ts

import * as assert from 'assert';
import * as fse from 'fs-extra';
import * as os from 'os';
import * as path from 'path';
import * as sinon from 'sinon';
import * as persistentState from '../../../common/persistentState';
import { getVenvForWorkspace, VENV_WORKSPACE_KEY } from '../../../managers/builtin/venvUtils';

/**
 * Helper that replicates the VIRTUAL_ENV subpath check from getVenvForWorkspace
 * using a specific path module, allowing cross-platform verification.
 */
function isVenvInsideWorkspace(
    fsPath: string,
    virtualEnv: string,
    pathModule: typeof path.posix | typeof path.win32,
): boolean {
    const rel = pathModule.relative(fsPath, virtualEnv);
    return rel === '' || (!rel.startsWith('..') && !pathModule.isAbsolute(rel));
}

suite('VIRTUAL_ENV subpath check - cross-platform', () => {
    suite('POSIX paths', () => {
        const p = path.posix;

        test('venv inside workspace should match', () => {
            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app/.venv', p), true);
        });

        test('venv deeply nested inside workspace should match', () => {
            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app/sub/dir/.venv', p), true);
        });

        test('venv equal to workspace should match', () => {
            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app', p), true);
        });

        test('sibling with shared prefix should NOT match', () => {
            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app2/.venv', p), false);
        });

        test('venv in completely different location should NOT match', () => {
            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/other/place/.venv', p), false);
        });

        test('parent directory should NOT match', () => {
            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/.venv', p), false);
        });
    });

    suite('Windows paths', () => {
        const p = path.win32;

        test('venv inside workspace should match', () => {
            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app\\.venv', p), true);
        });

        test('venv deeply nested inside workspace should match', () => {
            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app\\sub\\dir\\.venv', p), true);
        });

        test('venv equal to workspace should match', () => {
            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app', p), true);
        });

        test('sibling with shared prefix should NOT match', () => {
            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app2\\.venv', p), false);
        });

        test('venv on a different drive should NOT match', () => {
            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'D:\\proj\\app\\.venv', p), false);
        });

        test('parent directory should NOT match', () => {
            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\.venv', p), false);
        });

        test('case-insensitive drive letters should match', () => {
            // path.win32.relative handles case-insensitive drive letters
            assert.strictEqual(isVenvInsideWorkspace('c:\\proj\\app', 'C:\\proj\\app\\.venv', p), true);
        });

        test('UNC path inside workspace should match', () => {
            assert.strictEqual(
                isVenvInsideWorkspace('\\\\server\\share\\proj', '\\\\server\\share\\proj\\.venv', p),
                true,
            );
        });

        test('UNC path sibling should NOT match', () => {
            assert.strictEqual(
                isVenvInsideWorkspace('\\\\server\\share\\proj', '\\\\server\\share\\proj2\\.venv', p),
                false,
            );
        });
    });
});

suite('getVenvForWorkspace', () => {
    let mockState: {
        get: sinon.SinonStub;
        set: sinon.SinonStub;
        clear: sinon.SinonStub;
    };
    let getWorkspacePersistentStateStub: sinon.SinonStub;
    let originalVirtualEnv: string | undefined;
    let tmpDir: string;

    setup(async () => {
        originalVirtualEnv = process.env.VIRTUAL_ENV;
        delete process.env.VIRTUAL_ENV;

        tmpDir = path.join(os.tmpdir(), `venv-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
        await fse.ensureDir(tmpDir);

        mockState = {
            get: sinon.stub(),
            set: sinon.stub(),
            clear: sinon.stub(),
        };
        getWorkspacePersistentStateStub = sinon.stub(persistentState, 'getWorkspacePersistentState');
        getWorkspacePersistentStateStub.resolves(mockState);
    });

    teardown(async () => {
        if (originalVirtualEnv !== undefined) {
            process.env.VIRTUAL_ENV = originalVirtualEnv;
        } else {
            delete process.env.VIRTUAL_ENV;
        }
        sinon.restore();
        await fse.remove(tmpDir);
    });

    test('should return persisted selection when available', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        const venvPath = path.join(workspacePath, '.venv');
        await fse.ensureDir(venvPath);

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [workspacePath]: venvPath });

        const result = await getVenvForWorkspace(workspacePath);
        assert.strictEqual(result, venvPath, 'Should return persisted venv path');
    });

    test('should return persisted selection even when VIRTUAL_ENV is set', async () => {
        const projectA = path.join(tmpDir, 'projectA');
        const projectB = path.join(tmpDir, 'projectB');
        const persistedVenv = path.join(projectA, '.venv');
        const otherVenv = path.join(projectB, '.venv');
        await fse.ensureDir(persistedVenv);
        await fse.ensureDir(otherVenv);
        process.env.VIRTUAL_ENV = otherVenv;

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [projectA]: persistedVenv });

        const result = await getVenvForWorkspace(projectA);
        assert.strictEqual(result, persistedVenv, 'Persisted selection should take priority over VIRTUAL_ENV');
    });

    test('should fall back to VIRTUAL_ENV when no persisted selection and venv is inside workspace', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        const virtualEnvPath = path.join(workspacePath, '.venv');
        await fse.ensureDir(virtualEnvPath);
        process.env.VIRTUAL_ENV = virtualEnvPath;

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);

        const result = await getVenvForWorkspace(workspacePath);
        assert.strictEqual(result, virtualEnvPath, 'Should use VIRTUAL_ENV when it is inside the workspace');
    });

    test('should NOT use VIRTUAL_ENV when it points to a different project', async () => {
        const projectA = path.join(tmpDir, 'projectA');
        const projectB = path.join(tmpDir, 'projectB');
        const otherVenv = path.join(projectB, '.venv');
        await fse.ensureDir(projectA);
        await fse.ensureDir(otherVenv);
        process.env.VIRTUAL_ENV = otherVenv;

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);

        const result = await getVenvForWorkspace(projectA);
        assert.strictEqual(result, undefined, 'Should NOT use VIRTUAL_ENV from a different project');
    });

    test('should clear stale persisted path when venv no longer exists', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        const staleVenv = path.join(workspacePath, '.venv-old');
        await fse.ensureDir(workspacePath);
        // Note: staleVenv directory does NOT exist on disk

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [workspacePath]: staleVenv });

        const result = await getVenvForWorkspace(workspacePath);

        assert.strictEqual(result, undefined, 'Should return undefined for stale path');
        assert.ok(mockState.set.called, 'Should clear the stale entry from persistent state');
        const setArgs = mockState.set.firstCall.args;
        assert.strictEqual(setArgs[0], VENV_WORKSPACE_KEY, 'Should update the correct key');
        assert.strictEqual(setArgs[1][workspacePath], undefined, 'Should have removed the stale workspace entry');
    });

    test('should return undefined when no persisted selection and no VIRTUAL_ENV', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        await fse.ensureDir(workspacePath);

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);

        const result = await getVenvForWorkspace(workspacePath);
        assert.strictEqual(result, undefined, 'Should return undefined when nothing is available');
    });

    test('should return undefined when persisted data has no entry for this workspace', async () => {
        const projectA = path.join(tmpDir, 'projectA');
        const projectB = path.join(tmpDir, 'projectB');
        await fse.ensureDir(projectA);

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [projectB]: '/some/path' });

        const result = await getVenvForWorkspace(projectA);
        assert.strictEqual(result, undefined, 'Should return undefined when no entry for this workspace');
    });

    test('should fall back to VIRTUAL_ENV when data access throws inside try block', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        const virtualEnvPath = path.join(workspacePath, '.venv');
        await fse.ensureDir(virtualEnvPath);
        process.env.VIRTUAL_ENV = virtualEnvPath;

        // Return data object with a getter that throws when accessing the workspace key
        const badData: Record<string, string> = {};
        Object.defineProperty(badData, workspacePath, {
            get() {
                throw new Error('corrupted data');
            },
            enumerable: true,
            configurable: true,
        });
        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(badData);

        const result = await getVenvForWorkspace(workspacePath);
        assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when try block throws');
    });

    test('should not clear state when no envPath exists for the workspace key', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        await fse.ensureDir(workspacePath);

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ other: '/some/path' });

        await getVenvForWorkspace(workspacePath);

        assert.ok(!mockState.set.called, 'Should not call set when there is no stale entry to clear');
    });

    test('should fall back to VIRTUAL_ENV when state.get rejects', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        const virtualEnvPath = path.join(workspacePath, '.venv');
        await fse.ensureDir(virtualEnvPath);
        process.env.VIRTUAL_ENV = virtualEnvPath;

        mockState.get.withArgs(VENV_WORKSPACE_KEY).rejects(new Error('storage corrupted'));

        const result = await getVenvForWorkspace(workspacePath);
        assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when state.get rejects');
    });

    test('should fall back to VIRTUAL_ENV when getWorkspacePersistentState rejects', async () => {
        const workspacePath = path.join(tmpDir, 'projectA');
        const virtualEnvPath = path.join(workspacePath, '.venv');
        await fse.ensureDir(virtualEnvPath);
        process.env.VIRTUAL_ENV = virtualEnvPath;

        getWorkspacePersistentStateStub.rejects(new Error('persistent state unavailable'));

        const result = await getVenvForWorkspace(workspacePath);
        assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when persistent state fails');
    });

    test('should NOT use VIRTUAL_ENV for sibling path with shared prefix', async () => {
        const projectA = path.join(tmpDir, 'app');
        const siblingVenv = path.join(tmpDir, 'app2', '.venv');
        await fse.ensureDir(projectA);
        await fse.ensureDir(siblingVenv);
        process.env.VIRTUAL_ENV = siblingVenv;

        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);

        const result = await getVenvForWorkspace(projectA);
        assert.strictEqual(result, undefined, 'Should NOT match sibling path with shared prefix');
    });
});