ba-auth_keystore_keyvault/tests/KeyvaultKeys.test.ts

import { jest, describe, it, expect, beforeEach, afterEach } from '@jest/globals';
import type * as NodeJose from 'node-jose';
import type { GetKeyOptions, CreateKeyOptions, BeginDeleteKeyOptions, KeyVaultKey, PollerLike, PollOperationState, DeletedKey } from '@azure/keyvault-keys';


// --- Setup Mocks ---

// Mock @azure/identity
jest.unstable_mockModule('@azure/identity', () => ({
    ClientSecretCredential: jest.fn().mockImplementation(() => ({
        getToken: jest.fn()
    }))
}));

const sharedMockClient = {
    getKey: jest.fn<(name: string, options?: GetKeyOptions | undefined) => Promise<Partial<KeyVaultKey>>>(),
    createKey: jest.fn<(name: string, keyType: string, options?: CreateKeyOptions | undefined) => Promise<Partial<KeyVaultKey>>>(),
    beginDeleteKey: jest.fn<(name: string, options?: BeginDeleteKeyOptions | undefined) => Promise<PollerLike<PollOperationState<DeletedKey>, DeletedKey>>>(),
    listPropertiesOfKeys: jest.fn(),
};

// Mock @azure/keyvault-keys
jest.unstable_mockModule('@azure/keyvault-keys', () => ({
    // Configure the constructor to return this SHARED object every time
    KeyClient: jest.fn().mockReturnValue(sharedMockClient)
}));

// Mock Logging
jest.unstable_mockModule('@BridgemanAccessible/ba-logging', () => ({
    default: jest.fn(), // logMessage default export
    LogLevel: { DEBUG: 'debug', ERROR: 'error', WARN: 'warn' },
}));

// Mock BaseKeystore
jest.unstable_mockModule('@BridgemanAccessible/ba-auth_keystore', () => ({
    BaseKeystore: class BaseKeystore { keystore: any; },
    Keystore: { addKeystoreType: jest.fn() }
}));

// Mock node-forge (for fast key generation)
jest.unstable_mockModule('node-forge', () => ({
    pki: {
        rsa: {
            generateKeyPair: jest.fn().mockImplementation((bits, callback: any) => {
                // Return dummy key immediately
                callback(null, {
                    publicKey: { n: 'modulus', e: 'exponent' },
                    privateKey: { d: 'privateExponent' }
                });
            })
        }
    }
}));

// --- Import Dependencies ---
// Dynamic imports required due to unstable_mockModule
const { AzureKeyvaultCredentials } = await import('../src/KeyvaultCredentials.js');
const { KeyvaultKeys } = await import('../src/KeyvaultKeys.js');
const { KeyvaultKeystore } = await import('../src/KeyvaultKeystore.js');
const nodeJose = (await import('node-jose')).default;
const { JWK } = nodeJose;

// --- Helper for Async Iterators ---
// This mocks the `for await (const x of client.listPropertiesOfKeys())` syntax
const mockAsyncIterator = (items: any[]) => {
    return {
        [Symbol.asyncIterator]: async function* () {
            for (const item of items) {
                yield item;
            }
        }
    };
};

describe('Azure Key Vault Integration Tests', () => {
    const envVars = {
        KEY_VAULT_NAME: 'test-vault',
        KEY_VAULT_CLIENT_ID: 'cid',
        KEY_VAULT_CLIENT_SECRET: 'csec',
        KEY_VAULT_TENANT_ID: 'tid'
    };

    beforeEach(() => {
        jest.clearAllMocks();

        // Reset defaults for the shared mock
        sharedMockClient.listPropertiesOfKeys.mockReturnValue(mockAsyncIterator([]));
        sharedMockClient.createKey.mockResolvedValue({});
        sharedMockClient.beginDeleteKey.mockResolvedValue({} as any);
    });

    // --- Tests for Credentials ---
    describe('AzureKeyvaultCredentials', () => {
        it('should construct correct URL and credential objects', () => {
            const creds = new AzureKeyvaultCredentials('my-vault', 'id', 'secret', 'tenant');
            expect(creds.vaultUrl).toBe('https://my-vault.vault.azure.net');
            expect(creds.getCredentials()).toEqual({
                name: 'my-vault',
                client_id: 'id',
                client_secret: 'secret',
                tenant_id: 'tenant'
            });
        });
    });

    // --- Tests for KeyvaultKeys ---
    describe('KeyvaultKeys', () => {
        describe('init', () => {
            it('should throw if env vars are missing', async () => {
                await expect(KeyvaultKeys.init({ ...envVars, KEY_VAULT_NAME: '' }))
                    .rejects.toThrow('The KEY_VAULT_NAME variable must be set');
            });
        });

        describe('loadKeys', () => {
            it('should list keys, fetch details, and hydrate local state', async () => {
                // Mock List
                sharedMockClient.listPropertiesOfKeys.mockReturnValue(
                    mockAsyncIterator([
                        { name: 'key-1' }
                    ])
                );

                // Mock GetKey (Must return Azure structure)
                sharedMockClient.getKey.mockResolvedValue({
                    name: 'key-1',
                    key: {
                        kid: 'key-1',
                        kty: 'RSA',
                        n: new Uint8Array(new ArrayBuffer(1)),
                        e: new Uint8Array(new ArrayBuffer(1)),
                        keyOps: ['sign']
                    }
                });

                const kvKeys = await KeyvaultKeys.init(envVars);
                
                await kvKeys.loadKeys();

                // Verify interaction
                expect(sharedMockClient.listPropertiesOfKeys).toHaveBeenCalled();
                expect(sharedMockClient.getKey).toHaveBeenCalledWith('key-1');

                // Verify Local State
                const allKeys = kvKeys.all();
                expect(allKeys.length).toBe(1);
                expect(((allKeys as NodeJose.JWK.Key[])[0] as NodeJose.JWK.Key).kid).toBe('key-1');
                expect(((allKeys as NodeJose.JWK.Key[])[0] as NodeJose.JWK.Key).alg).toBe('RS256'); // Verifies conversion logic ran
            });

            it('should handle empty vault gracefully', async () => {
                sharedMockClient.listPropertiesOfKeys.mockReturnValue(mockAsyncIterator([]));
                
                const kvKeys = await KeyvaultKeys.init(envVars);
                
                const keys = await kvKeys.loadKeys();
                expect(keys).toEqual([]);
                expect(kvKeys.all()).toEqual([]);
            });
        });

        describe('add', () => {
            it('should write to Azure AND update local state', async () => {
                const kvKeys = await KeyvaultKeys.init(envVars);
                
                sharedMockClient.createKey.mockResolvedValue({}); // Return value doesn't matter much for this, but could

                const rawKey = { 
                    kty: 'oct', 
                    k: '1234', 
                    kid: 'new-key',
                    use: 'sig' 
                };
                
                await kvKeys.add(rawKey as any);

                // Check Azure Write
                expect(sharedMockClient.createKey).toHaveBeenCalledWith(
                    'new-key',
                    'oct',
                    expect.objectContaining({ keyOps: ['sign', 'verify'] })
                );

                // Check Local State (THIS WILL FAIL ON CURRENT IMPLEMENTATION)
                const storedKey = kvKeys.get('new-key');
                expect(storedKey).toBeDefined();
                expect(storedKey.kid).toBe('new-key');
            });
        });

        describe('generate', () => {
            it('should generate key via node-forge and save to Azure/Local', async () => {
                const kvKeys = await KeyvaultKeys.init(envVars);
                
                // Mock createKey to return success
                sharedMockClient.createKey.mockResolvedValue({});

                const generated = await kvKeys.generate('RSA', 2048, { use: 'sig' });

                // Verify Azure call
                expect(sharedMockClient.createKey).toHaveBeenCalled();
                const args = sharedMockClient.createKey.mock.calls[0];
                expect((args as [name: string, keyType: string, options?: CreateKeyOptions | undefined])[0]).toBe(generated.kid); // KID match
                
                // Verify Local State
                expect(kvKeys.get(generated.kid)).toBeDefined();
            });
        });

        describe('remove', () => {
            it('should delete from Azure and remove from local state', async () => {
                const kvKeys = await KeyvaultKeys.init(envVars);
                
                // Add a dummy key to local state
                const dummyKey = await JWK.createKey('oct', 256, { kid: 'del-key' });
                (kvKeys as any).keys.push(dummyKey);

                // Mock delete promise
                sharedMockClient.beginDeleteKey.mockResolvedValue({} as any);

                kvKeys.remove(dummyKey);

                // Verify Azure call
                expect(sharedMockClient.beginDeleteKey).toHaveBeenCalledWith('del-key');

                // Verify Local removal
                expect(() => kvKeys.get('del-key')).toThrow();
            });
        });
    });

    // --- Tests for KeyvaultKeystore ---
    describe('KeyvaultKeystore', () => {
        it('should hydrate keys on init', async () => {
            // Setup mocking for the static init sequence
            sharedMockClient.listPropertiesOfKeys.mockReturnValue(mockAsyncIterator([{ name: 'existing-k' }]));
            sharedMockClient.getKey.mockResolvedValue({ 
                name: 'existing-k',
                key: { kid: 'existing-k', kty: 'RSA' },
                properties: { 
                    name: 'existing-k',
                    vaultUrl: 'https://test-vault.vault.azure.net',
                    enabled: true 
                }
            });

            // We need to hijack the constructor internally used by init
            // Since we mocked the module, the classes behave as defined in src,
            // but we need to ensure the KeyClient instance used inside is our mock.
            // The easiest way with `unstable_mockModule` in integration tests 
            // is relying on the fact that `new KeyClient()` returns our mock.
            
            const keystore = await KeyvaultKeystore.init(envVars);

            expect(sharedMockClient.listPropertiesOfKeys).toHaveBeenCalled();
            expect(sharedMockClient.getKey).toHaveBeenCalledWith('existing-k');
        });

        it('should setup initial keys if vault is empty', async () => {
            sharedMockClient.listPropertiesOfKeys.mockReturnValue(mockAsyncIterator([]));
            sharedMockClient.createKey.mockResolvedValue({}); // Allow generate to succeed

            await KeyvaultKeystore.init(envVars);

            // Should trigger generate -> createKey twice (sig + enc)
            expect(sharedMockClient.createKey).toHaveBeenCalledTimes(2);
        });
    });
});