ba-auth_keystore_vault/tests/VaultKeyst.test.ts

import { jest, describe, it, expect, beforeEach, afterEach } from '@jest/globals';
import type * as NodeJose from 'node-jose';

import type vault from 'node-vault';

import type { VaultCredentials as VaultCredentialsType } from '../src/VaultCredentials.js';
import type { VaultKeys as VaultKeysType } from '../src/VaultKeys.js';
import type { VaultKeystore as VaultKeystoreType } from '../src/VaultKeystore.js';

type JestVaultClient = Partial<vault.client> & {
    approleLogin: jest.Mock<(options?: vault.Option | undefined) => Promise<any>>,
    write: jest.Mock<(path: string, data: any, requestOptions?: vault.Option | undefined) => Promise<any>>,
    read: jest.Mock<(path: string, requestOptions?: vault.Option | undefined) => Promise<any>>,
    list: jest.Mock<(path: string, requestOptions?: vault.Option | undefined) => Promise<any>>,
    token: string
};

// Mocks
jest.unstable_mockModule('node-vault', () => ({
    'default': jest.fn<(options?: vault.VaultOptions | undefined) => JestVaultClient>().mockReturnValue({
        approleLogin: jest.fn<(options?: vault.Option | undefined) => Promise<any>>().mockResolvedValue({ auth: { client_token: 'fake-token' } }),
        write: jest.fn<(path: string, data: any, requestOptions?: vault.Option | undefined) => Promise<any>>().mockResolvedValue({}),
        read: jest.fn<(path: string, requestOptions?: vault.Option | undefined) => Promise<any>>(),
        list: jest.fn<(path: string, requestOptions?: vault.Option | undefined) => Promise<any>>(),
        token: ''
    })
}));

jest.unstable_mockModule('@BridgemanAccessible/ba-logging', () => ({
    logMessage: jest.fn<(message: string, level: 'debug' | 'error' | 'warn') => void>(),
    LogLevel: { DEBUG: 'debug', ERROR: 'error', WARN: 'warn' },
}));

// Mock the BaseKeystore dependency to avoid importing logic we aren't testing
jest.unstable_mockModule('@BridgemanAccessible/ba-auth_keystore', () => ({
  BaseKeystore: class BaseKeystore { keystore: any; },
  Keystore: { addKeystoreType: jest.fn() }
}));

jest.unstable_mockModule('node-forge', () => ({
    pki: {
        rsa: {
            generateKeyPair: jest.fn<(bits: number, callback: (err: any, keypair: { publicKey: any, privateKey: any }) => void) => void>().mockImplementation((bits, callback) => {
                // Simulate async key generation
                setTimeout(() => {
                    callback(null, {
                        publicKey: { n: 'modulus', e: 'exponent' },
                        privateKey: { d: 'privateExponent' }
                    });
                }, 10);
            })
        }
    }
}))

describe('Vault Integration Tests', () => {
    // Mock Vault Client Instance
    let nodeVault: jest.Mock<(options?: vault.VaultOptions | undefined) => JestVaultClient>;
    let mockVaultClient: JestVaultClient;

    let JWK: typeof NodeJose.JWK;

    let VaultKeys: typeof VaultKeysType;
    let VaultCredentials: typeof VaultCredentialsType;
    let VaultKeystore: typeof VaultKeystoreType;
    
    // Standard Test Env Vars
    const envVars = {
        VAULT_NAME: 'test-vault',
        VAULT_ROLE_ID: 'role-id',
        VAULT_SECRET_ID: 'secret-id',
        VAULT_PORT: 8200,
    };
    
    beforeEach(async () => {
        // Reset mocks
        jest.clearAllMocks();
        
        // Setup the Node-Vault mock return values
        const nodeVaultModule = await import('node-vault');
        nodeVault = nodeVaultModule.default as jest.Mock<(options?: vault.VaultOptions | undefined) => JestVaultClient>;
        mockVaultClient = nodeVault();

        const nodeJoseModule = await import('node-jose');
        JWK = nodeJoseModule.default.JWK;

        // Setup the VaultCredentials and VaultKeys imports (after mocks)
        const vaultCredentialsModule = await import('../src/VaultCredentials.js');
        VaultCredentials = vaultCredentialsModule.VaultCredentials;
        
        const vaultKeysModule = await import('../src/VaultKeys.js');
        VaultKeys = vaultKeysModule.VaultKeys;

        const vaultKeystoreModule = await import('../src/VaultKeystore.js');
        VaultKeystore = vaultKeystoreModule.VaultKeystore;
    });
    
    // --- 1. Testing VaultCredentials ---
    describe('VaultCredentials', () => {
        it('should initialize and perform appRole login', async () => {
            // Call the init method
            const creds = await VaultCredentials.init('localhost', 'role', 'secret', 8200);
            
            // Check the vault client was created correctly
            expect(nodeVault).toHaveBeenCalledWith(expect.objectContaining({ endpoint: 'http://localhost:8200' }));
            expect(mockVaultClient.approleLogin).toHaveBeenCalledWith({ role_id: 'role', secret_id: 'secret' });
            expect(creds.getVaultClient().token).toBe('fake-token');
        });
    });
    
    // --- 2. Testing VaultKeys ---
    describe('VaultKeys', () => {
        it('should throw error if env vars are missing', async () => {
            await expect(VaultKeys.init({ ...envVars, VAULT_NAME: '' })).rejects.toThrow('The VAULT_NAME variable must be set');
        });
        
        it('should initialize with empty local keys', async () => {
            const vaultKeys = await VaultKeys.init(envVars);
            
            // "all" is synchronous and reads from the local private array
            expect(vaultKeys.all()).toEqual([]);
        });
        
        describe('loadKeys (Hydration)', () => {
            it('should load keys from Vault and populate local state', async () => {
                // 1. Setup a real JWK to return from the mock
                const realKey = await JWK.createKey('oct', 256, { kid: 'test-key-1' });
                const keyJson = realKey.toJSON(true);
                
                // 2. Mock list() to return one key ID
                mockVaultClient.list.mockResolvedValue({ data: { keys: ['test-key-1'] } });
                
                // 3. Mock read() to return the key data
                mockVaultClient.read.mockResolvedValue({ data: { data: keyJson } });
                
                const vaultKeys = await VaultKeys.init(envVars);
                await vaultKeys.loadKeys();
                
                // 4. Verify list was called
                expect(mockVaultClient.list).toHaveBeenCalledWith('secret/data/keys');
                
                // 5. Verify read was called
                expect(mockVaultClient.read).toHaveBeenCalledWith('secret/data/keys/test-key-1');
                
                // 6. Verify local state is updated (The synchronous get/all should work now)
                const keys = vaultKeys.all();
                expect(keys.length).toBe(1);
                expect(keys[0]?.kid).toBe('test-key-1');
            });
            
            it('should handle empty vault gracefully (404 on list)', async () => {
                // Simulate a 404 from Vault (standard behavior when path doesn't exist yet)
                const err: any = new Error('Not Found');
                err.response = { statusCode: 404 };
                mockVaultClient.list.mockRejectedValue(err);
                
                const vaultKeys = await VaultKeys.init(envVars);
                const keys = await vaultKeys.loadKeys();
                
                expect(keys).toEqual([]);
                expect(vaultKeys.all()).toEqual([]);
            });
        });
        
        describe('add', () => {
            it('should write to Vault AND update local state', async () => {
                const vaultKeys = await VaultKeys.init(envVars);
                
                // Input: Raw Key Data
                const rawKey = { kty: 'oct', k: '1234', kid: 'new-key' };
                
                // Call add
                await vaultKeys.add(rawKey as any);
                
                // 1. Check Vault Write
                expect(mockVaultClient.write).toHaveBeenCalledWith(
                    'secret/data/keys/new-key',
                    expect.objectContaining({ 
                        data: expect.objectContaining({ 
                            kid: 'new-key' 
                        }) 
                    })
                );
                
                // 2. Check Local State
                const storedKey = vaultKeys.get('new-key');
                expect(storedKey).toBeDefined();
                expect(storedKey.kid).toBe('new-key');
            });
        });
        
        describe('generate', () => {
            it('should generate a key, write it to vault, and store it locally', async () => {
                const vaultKeys = await VaultKeys.init(envVars);
                
                // This uses the real node-jose generator logic
                const generatedKey = await vaultKeys.generate('RSA', 2048, { use: 'sig' });
                
                // Verify Vault interaction
                expect(mockVaultClient.write).toHaveBeenCalled();
                const writeCall = mockVaultClient.write.mock.calls[0];
                // Ensure the path contains the generated KID
                expect((writeCall as [path: string, data: any, requestOptions?: vault.Option | undefined])[0]).toContain(`secret/data/keys/${generatedKey.kid}`);
                
                // Verify Local interaction
                expect(vaultKeys.get(generatedKey.kid)).toBeDefined();
                expect(generatedKey.kty).toBe('RSA');
            });
        });
        
        describe('get (Synchronous)', () => {
            it('should retrieve a key by kid from local state', async () => {
                const vaultKeys = await VaultKeys.init(envVars);
                
                // Manually inject a key via add() first to populate state
                const addedKey = await vaultKeys.add(await JWK.createKey('oct', 256, {}));
                
                // Retrieve
                const retrieved = vaultKeys.get(addedKey.kid);
                expect(retrieved).toBeDefined();
                expect(retrieved.kid).toBe(addedKey.kid);
            });
            
            it('should throw if key is missing locally', async () => {
                const vaultKeys = await VaultKeys.init(envVars);
                expect(() => vaultKeys.get('non-existent')).toThrow('Key non-existent not found');
            });
        });
    });
    
    // --- 3. Testing VaultKeystore (Wrapper) ---
    describe('VaultKeystore', () => {
        it('should initialize keys and setup initial keys if vault is empty', async () => {
            // Mock empty vault
            mockVaultClient.list.mockRejectedValue({ response: { statusCode: 404 } });
            
            const keystore = await VaultKeystore.init(envVars);
            
            // It should have tried to list keys
            expect(mockVaultClient.list).toHaveBeenCalled();
            
            // Since list failed (empty), it calls setupInitialKeys -> generate -> write
            // We expect at least 2 writes (signing key and encryption key)
            expect(mockVaultClient.write).toHaveBeenCalledTimes(2);
        });
        
        it('should load existing keys and NOT generate new ones', async () => {
            // Mock existing keys
            mockVaultClient.list.mockResolvedValue({ data: { keys: ['existing-key'] } });
            const realKey = await JWK.createKey('RSA', 2048, { kid: 'existing-key' });
            mockVaultClient.read.mockResolvedValue({ data: { data: realKey.toJSON(true) } });
            
            await VaultKeystore.init(envVars);
            
            // Should read the key
            expect(mockVaultClient.read).toHaveBeenCalled();
            // Should NOT write new keys (generate)
            expect(mockVaultClient.write).not.toHaveBeenCalled();
        });
    });
});