sentryagent-idp/tests/unit/services/DIDService.test.ts

/**
 * Unit tests for src/services/DIDService.ts
 * All public methods are tested. pg Pool, node-vault, and Redis are mocked.
 */

import { DIDService } from '../../../src/services/DIDService';
import { AgentNotFoundError } from '../../../src/utils/errors';
import { IPublicKeyJwk } from '../../../src/types/did';

// ─── Mock node-vault ─────────────────────────────────────────────────────────

jest.mock('node-vault', () => {
  return jest.fn(() => ({
    write: jest.fn().mockResolvedValue({}),
    read: jest.fn().mockResolvedValue({}),
  }));
});

// ─── Mock pg Pool ─────────────────────────────────────────────────────────────

const mockQuery = jest.fn();
const mockPool = {
  query: mockQuery,
} as never;

// ─── Mock Redis client ────────────────────────────────────────────────────────

const mockRedisGet = jest.fn();
const mockRedisSet = jest.fn();
const mockRedis = {
  get: mockRedisGet,
  set: mockRedisSet,
} as never;

// ─── Constants ────────────────────────────────────────────────────────────────

const TEST_DOMAIN = 'test.example.com';
const AGENT_ID = 'agt_test';
const ORG_ID = 'org_system';

const MOCK_PUBLIC_KEY_JWK: IPublicKeyJwk = {
  kty: 'EC',
  crv: 'P-256',
  x: 'abc123',
  y: 'def456',
};

const MOCK_AGENT_ROW = {
  agent_id: AGENT_ID,
  organization_id: ORG_ID,
  email: 'test@example.com',
  agent_type: 'orchestrator',
  version: '1.0.0',
  capabilities: ['task-planning'],
  owner: 'acme',
  deployment_env: 'production',
  status: 'active',
  created_at: new Date('2026-01-01T00:00:00Z'),
  updated_at: new Date('2026-01-02T00:00:00Z'),
  did: `did:web:${TEST_DOMAIN}:agents:${AGENT_ID}`,
  did_created_at: new Date('2026-01-01T00:00:00Z'),
  key_id: 'key_test',
  public_key_jwk: MOCK_PUBLIC_KEY_JWK,
  vault_key_path: 'dev:no-vault',
  key_type: 'EC',
  curve: 'P-256',
  key_created_at: new Date('2026-01-01T00:00:00Z'),
};

const MOCK_DECOMMISSIONED_ROW = {
  ...MOCK_AGENT_ROW,
  status: 'decommissioned',
};

/** Row simulating an agent that has no key yet (LEFT JOIN returns NULLs for key columns). */
const MOCK_AGENT_ROW_NO_KEY = {
  agent_id: AGENT_ID,
  organization_id: ORG_ID,
  email: 'test@example.com',
  agent_type: 'orchestrator',
  version: '1.0.0',
  capabilities: ['task-planning'],
  owner: 'acme',
  deployment_env: 'production',
  status: 'active',
  created_at: new Date('2026-01-01T00:00:00Z'),
  updated_at: new Date('2026-01-02T00:00:00Z'),
  did: null,
  did_created_at: null,
  key_id: null,
  public_key_jwk: null,
  vault_key_path: null,
  key_type: null,
  curve: null,
  key_created_at: null,
};

// ─── Helper ───────────────────────────────────────────────────────────────────

function buildService(): DIDService {
  return new DIDService(mockPool, null, mockRedis);
}

/** Reset all mocks to clean state before each test. */
function resetMocks(): void {
  jest.clearAllMocks();
  mockRedisGet.mockResolvedValue(null);
  mockRedisSet.mockResolvedValue('OK');
  // Default pool.query returns empty result (agent not found)
  mockQuery.mockResolvedValue({ rows: [] });
}

// ─── Suite ────────────────────────────────────────────────────────────────────

describe('DIDService', () => {
  beforeAll(() => {
    process.env['DID_WEB_DOMAIN'] = TEST_DOMAIN;
  });

  afterAll(() => {
    delete process.env['DID_WEB_DOMAIN'];
    delete process.env['VAULT_ADDR'];
    delete process.env['VAULT_TOKEN'];
  });

  beforeEach(() => {
    resetMocks();
  });

  // ─── generateDIDForAgent ──────────────────────────────────────────────────

  describe('generateDIDForAgent()', () => {
    it('should return a did:web DID and public key JWK', async () => {
      mockQuery
        .mockResolvedValueOnce({ rows: [] }) // INSERT agent_did_keys
        .mockResolvedValueOnce({ rows: [] }); // UPDATE agents

      const service = buildService();
      const { did, publicKeyJwk } = await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      expect(did).toBe(`did:web:${TEST_DOMAIN}:agents:${AGENT_ID}`);
      expect(publicKeyJwk).toBeDefined();
      expect(publicKeyJwk.kty).toBe('EC');
      expect(publicKeyJwk.crv).toBe('P-256');
    });

    it('should call pool.query twice (INSERT + UPDATE)', async () => {
      mockQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [] });

      const service = buildService();
      await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      expect(mockQuery).toHaveBeenCalledTimes(2);
    });

    it('should INSERT into agent_did_keys with correct columns', async () => {
      mockQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [] });

      const service = buildService();
      await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      const insertCall = mockQuery.mock.calls[0] as [string, unknown[]];
      expect(insertCall[0]).toContain('INSERT INTO agent_did_keys');
      expect(insertCall[1]).toContain(AGENT_ID);
      expect(insertCall[1]).toContain(ORG_ID);
    });

    it('should UPDATE agents table with the generated DID', async () => {
      mockQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [] });

      const service = buildService();
      const { did } = await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      const updateCall = mockQuery.mock.calls[1] as [string, unknown[]];
      expect(updateCall[0]).toContain('UPDATE agents');
      expect(updateCall[1]).toContain(did);
      expect(updateCall[1]).toContain(AGENT_ID);
    });

    it('should use dev:no-vault marker when Vault env vars are not set', async () => {
      delete process.env['VAULT_ADDR'];
      delete process.env['VAULT_TOKEN'];
      mockQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [] });

      const service = buildService();
      await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      const insertCall = mockQuery.mock.calls[0] as [string, unknown[]];
      // vault_key_path is the 5th param ($5)
      expect(insertCall[1][4]).toBe('dev:no-vault');
    });

    it('should call vault.write when VAULT_ADDR and VAULT_TOKEN are set', async () => {
      process.env['VAULT_ADDR'] = 'http://localhost:8200';
      process.env['VAULT_TOKEN'] = 'test-token';

      const nodeVault = require('node-vault');
      const mockVaultInstance = { write: jest.fn().mockResolvedValue({}) };
      (nodeVault as jest.Mock).mockReturnValueOnce(mockVaultInstance);

      mockQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [] });

      const service = buildService();
      await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      expect(mockVaultInstance.write).toHaveBeenCalledTimes(1);
      const [vaultPath, payload] = mockVaultInstance.write.mock.calls[0] as [
        string,
        { data: { privateKeyPem: string } },
      ];
      expect(vaultPath).toContain(AGENT_ID);
      expect(payload.data.privateKeyPem).toBeDefined();
      expect(payload.data.privateKeyPem).toContain('-----BEGIN');

      delete process.env['VAULT_ADDR'];
      delete process.env['VAULT_TOKEN'];
    });

    it('should NEVER include the private key PEM in the returned object', async () => {
      mockQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [] });

      const service = buildService();
      const result = await service.generateDIDForAgent(AGENT_ID, ORG_ID);

      const serialised = JSON.stringify(result);
      expect(serialised).not.toContain('privateKeyPem');
      expect(serialised).not.toContain('PRIVATE KEY');
    });
  });

  // ─── buildInstanceDIDDocument ─────────────────────────────────────────────

  describe('buildInstanceDIDDocument()', () => {
    it('should return a W3C DID Document with correct @context', async () => {
      const service = buildService();
      const doc = await service.buildInstanceDIDDocument();

      expect(doc['@context']).toContain('https://www.w3.org/ns/did/v1');
    });

    it('should set the DID id to did:web:{domain}', async () => {
      const service = buildService();
      const doc = await service.buildInstanceDIDDocument();

      expect(doc.id).toBe(`did:web:${TEST_DOMAIN}`);
    });

    it('should include a verificationMethod of type JsonWebKey2020', async () => {
      const service = buildService();
      const doc = await service.buildInstanceDIDDocument();

      expect(doc.verificationMethod).toHaveLength(1);
      expect(doc.verificationMethod[0].type).toBe('JsonWebKey2020');
      expect(doc.verificationMethod[0].controller).toBe(`did:web:${TEST_DOMAIN}`);
    });

    it('should include an AgentIdentityProvider service endpoint', async () => {
      const service = buildService();
      const doc = await service.buildInstanceDIDDocument();

      expect(doc.service).toBeDefined();
      const svc = doc.service![0];
      expect(svc.type).toBe('AgentIdentityProvider');
      expect(svc.serviceEndpoint).toContain(TEST_DOMAIN);
    });

    it('should cache the document in Redis on first call', async () => {
      const service = buildService();
      await service.buildInstanceDIDDocument();

      expect(mockRedisSet).toHaveBeenCalledWith(
        'did:doc:instance',
        expect.any(String),
        { EX: expect.any(Number) },
      );
    });

    it('should return cached document on second call without storing again', async () => {
      const service = buildService();

      // First call — cache miss, builds and stores
      const doc = await service.buildInstanceDIDDocument();

      // Reset and simulate cache hit on second call
      resetMocks();
      mockRedisGet.mockResolvedValueOnce(JSON.stringify(doc));

      await service.buildInstanceDIDDocument();

      // Only called once total for the second call
      expect(mockRedisGet).toHaveBeenCalledTimes(1);
      // set should NOT be called on cache hit
      expect(mockRedisSet).not.toHaveBeenCalled();
    });

    it('should throw an error when DID_WEB_DOMAIN is not set', async () => {
      delete process.env['DID_WEB_DOMAIN'];
      const service = buildService();
      await expect(service.buildInstanceDIDDocument()).rejects.toThrow(
        'DID_WEB_DOMAIN environment variable is required',
      );
      process.env['DID_WEB_DOMAIN'] = TEST_DOMAIN;
    });
  });

  // ─── buildAgentDIDDocument ────────────────────────────────────────────────

  describe('buildAgentDIDDocument()', () => {
    it('should return a DID Document for an active agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const result = await service.buildAgentDIDDocument(AGENT_ID);

      expect(result.deactivated).toBe(false);
      expect(result.document.id).toBe(`did:web:${TEST_DOMAIN}:agents:${AGENT_ID}`);
    });

    it('should include DID contexts for active agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const { document } = await service.buildAgentDIDDocument(AGENT_ID);

      expect(document['@context']).toContain('https://www.w3.org/ns/did/v1');
      expect(document['@context']).toContain('https://w3id.org/agntcy/v1');
    });

    it('should include the public key in verificationMethod', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const { document } = await service.buildAgentDIDDocument(AGENT_ID);

      expect(document.verificationMethod).toHaveLength(1);
      expect(document.verificationMethod[0].publicKeyJwk).toEqual(MOCK_PUBLIC_KEY_JWK);
    });

    it('should include AGNTCY extension with agent fields', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const { document } = await service.buildAgentDIDDocument(AGENT_ID);

      expect(document.agntcy).toBeDefined();
      expect(document.agntcy!.agentId).toBe(AGENT_ID);
      expect(document.agntcy!.agentType).toBe('orchestrator');
      expect(document.agntcy!.capabilities).toEqual(['task-planning']);
    });

    it('should return deactivated=true for a decommissioned agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_DECOMMISSIONED_ROW] });

      const service = buildService();
      const result = await service.buildAgentDIDDocument(AGENT_ID);

      expect(result.deactivated).toBe(true);
    });

    it('should include AgentStatus service endpoint for decommissioned agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_DECOMMISSIONED_ROW] });

      const service = buildService();
      const { document } = await service.buildAgentDIDDocument(AGENT_ID);

      expect(document.service).toBeDefined();
      const statusSvc = document.service!.find((s) => s.type === 'AgentStatus');
      expect(statusSvc).toBeDefined();
      expect(statusSvc!.serviceEndpoint).toBe('decommissioned');
    });

    it('should NOT cache document for decommissioned agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_DECOMMISSIONED_ROW] });

      const service = buildService();
      await service.buildAgentDIDDocument(AGENT_ID);

      expect(mockRedisSet).not.toHaveBeenCalled();
    });

    it('should cache document for active agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      await service.buildAgentDIDDocument(AGENT_ID);

      expect(mockRedisSet).toHaveBeenCalledWith(
        `did:doc:${AGENT_ID}`,
        expect.any(String),
        { EX: expect.any(Number) },
      );
    });

    it('should use cached document on second call for active agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const { document: firstDoc } = await service.buildAgentDIDDocument(AGENT_ID);

      // Reset and simulate cache hit for second call
      resetMocks();
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });
      mockRedisGet.mockResolvedValueOnce(JSON.stringify(firstDoc));

      const { document: secondDoc } = await service.buildAgentDIDDocument(AGENT_ID);

      expect(secondDoc).toEqual(firstDoc);
      // set should NOT be called on cache hit
      expect(mockRedisSet).not.toHaveBeenCalled();
    });

    it('should throw AgentNotFoundError when agent does not exist', async () => {
      // Default mock already returns { rows: [] } from resetMocks
      const service = buildService();
      await expect(service.buildAgentDIDDocument('nonexistent-id')).rejects.toThrow(
        AgentNotFoundError,
      );
    });

    it('should handle agent with no key (empty verificationMethod)', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW_NO_KEY] });

      const service = buildService();
      const { document } = await service.buildAgentDIDDocument(AGENT_ID);

      expect(document.verificationMethod).toHaveLength(0);
      expect(document.authentication).toHaveLength(0);
      expect(document.assertionMethod).toHaveLength(0);
    });

    it('should NEVER include private key material in the DID Document', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const { document } = await service.buildAgentDIDDocument(AGENT_ID);

      const serialised = JSON.stringify(document);
      expect(serialised).not.toContain('privateKeyPem');
      expect(serialised).not.toContain('PRIVATE KEY');
    });
  });

  // ─── buildResolutionResult ────────────────────────────────────────────────

  describe('buildResolutionResult()', () => {
    it('should return a W3C DID Resolution result with all required sections', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const result = await service.buildResolutionResult(AGENT_ID);

      expect(result.didDocument).toBeDefined();
      expect(result.didDocumentMetadata).toBeDefined();
      expect(result.didResolutionMetadata).toBeDefined();
    });

    it('should set deactivated=false in metadata for active agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const result = await service.buildResolutionResult(AGENT_ID);

      expect(result.didDocumentMetadata.deactivated).toBe(false);
    });

    it('should set deactivated=true in metadata for decommissioned agent', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_DECOMMISSIONED_ROW] });

      const service = buildService();
      const result = await service.buildResolutionResult(AGENT_ID);

      expect(result.didDocumentMetadata.deactivated).toBe(true);
    });

    it('should set contentType to application/did+ld+json in resolutionMetadata', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const result = await service.buildResolutionResult(AGENT_ID);

      expect(result.didResolutionMetadata.contentType).toBe('application/did+ld+json');
    });

    it('should include ISO timestamps for created, updated, and retrieved', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const result = await service.buildResolutionResult(AGENT_ID);

      expect(() => new Date(result.didDocumentMetadata.created)).not.toThrow();
      expect(() => new Date(result.didDocumentMetadata.updated)).not.toThrow();
      expect(() => new Date(result.didResolutionMetadata.retrieved)).not.toThrow();
    });

    it('should throw AgentNotFoundError when agent does not exist', async () => {
      // Default mock already returns { rows: [] } from resetMocks
      const service = buildService();
      await expect(service.buildResolutionResult('nonexistent-id')).rejects.toThrow(
        AgentNotFoundError,
      );
    });

    it('should NEVER include private key material in the resolution result', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const result = await service.buildResolutionResult(AGENT_ID);

      const serialised = JSON.stringify(result);
      expect(serialised).not.toContain('privateKeyPem');
      expect(serialised).not.toContain('PRIVATE KEY');
    });
  });

  // ─── buildAgentCard ───────────────────────────────────────────────────────

  describe('buildAgentCard()', () => {
    it('should return an AGNTCY agent card with correct fields', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const card = await service.buildAgentCard(AGENT_ID);

      expect(card.did).toBe(`did:web:${TEST_DOMAIN}:agents:${AGENT_ID}`);
      expect(card.name).toBe('test@example.com');
      expect(card.agentType).toBe('orchestrator');
      expect(card.capabilities).toEqual(['task-planning']);
      expect(card.owner).toBe('acme');
      expect(card.version).toBe('1.0.0');
      expect(card.deploymentEnv).toBe('production');
      expect(card.identityProvider).toBe('https://idp.sentryagent.ai');
    });

    it('should include a valid ISO issuedAt timestamp', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const card = await service.buildAgentCard(AGENT_ID);

      expect(() => new Date(card.issuedAt)).not.toThrow();
      expect(new Date(card.issuedAt).toISOString()).toBe(card.issuedAt);
    });

    it('should use agent.createdAt as issuedAt when key_created_at is null', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW_NO_KEY] });

      const service = buildService();
      const card = await service.buildAgentCard(AGENT_ID);

      expect(card.issuedAt).toBe(MOCK_AGENT_ROW_NO_KEY.created_at.toISOString());
    });

    it('should use key_created_at as issuedAt when present', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const card = await service.buildAgentCard(AGENT_ID);

      expect(card.issuedAt).toBe(MOCK_AGENT_ROW.key_created_at.toISOString());
    });

    it('should throw AgentNotFoundError when agent does not exist', async () => {
      // Default mock already returns { rows: [] } from resetMocks
      const service = buildService();
      await expect(service.buildAgentCard('nonexistent-id')).rejects.toThrow(AgentNotFoundError);
    });

    it('should NEVER include private key material in the agent card', async () => {
      mockQuery.mockResolvedValueOnce({ rows: [MOCK_AGENT_ROW] });

      const service = buildService();
      const card = await service.buildAgentCard(AGENT_ID);

      const serialised = JSON.stringify(card);
      expect(serialised).not.toContain('privateKeyPem');
      expect(serialised).not.toContain('PRIVATE KEY');
    });
  });
});