add unit tests

dydxprotocol · Mar 1, 2024 · 2ae9e5c · 2ae9e5c
1 parent e433ca3
commit 2ae9e5c
Showing 1 changed file with 175 additions and 4 deletions.
diff --git a/indexer/services/comlink/__tests__/controllers/api/v4/compliance-v2-controller.test.ts b/indexer/services/comlink/__tests__/controllers/api/v4/compliance-v2-controller.test.ts
@@ -1,11 +1,11 @@
 import {
+  ComplianceReason,
   ComplianceStatus,
+  ComplianceStatusFromDatabase,
+  ComplianceStatusTable,
   dbHelpers,
-  testMocks,
   testConstants,
-  ComplianceReason,
-  ComplianceStatusTable,
-  ComplianceStatusFromDatabase,
+  testMocks,
 } from '@dydxprotocol-indexer/postgres';
 import { getIpAddr } from '../../../../src/lib/utils';
 import { sendRequest } from '../../../helpers/helpers';
@@ -18,6 +18,11 @@ import config from '../../../../src/config';
 import { DateTime } from 'luxon';
 import { ComplianceAction } from '../../../../src/controllers/api/v4/compliance-v2-controller';
 import { ExtendedSecp256k1Signature, Secp256k1, sha256 } from '@cosmjs/crypto';
+import { getGeoComplianceReason } from '../../../../src/helpers/compliance/compliance-utils';
+import { isRestrictedCountryHeaders } from '@dydxprotocol-indexer/compliance';
+
+jest.mock('@dydxprotocol-indexer/compliance');
+jest.mock('../../../../src/helpers/compliance/compliance-utils');
 
 jest.mock('../../../../src/lib/utils', () => ({
   ...jest.requireActual('../../../../src/lib/utils'),
@@ -231,7 +236,12 @@ describe('ComplianceV2Controller', () => {
   });
 
   describe('POST /geoblock', () => {
+    let getGeoComplianceReasonSpy: jest.SpyInstance;
+    let isRestrictedCountryHeadersSpy: jest.SpyInstance;
+
     beforeEach(async () => {
+      getGeoComplianceReasonSpy = getGeoComplianceReason as unknown as jest.Mock;
+      isRestrictedCountryHeadersSpy = isRestrictedCountryHeaders as unknown as jest.Mock;
       ipAddrMock.mockReturnValue(ipAddr);
       await testMocks.seedData();
       jest.mock('@cosmjs/crypto', () => ({
@@ -322,5 +332,166 @@ describe('ComplianceV2Controller', () => {
       expect(response.status).toEqual(200);
       expect(response.body.status).toEqual(ComplianceStatus.COMPLIANT);
     });
+
+    it('should set status to BLOCKED for ONBOARD action from a restricted country with no existing compliance status', async () => {
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          address: testConstants.defaultAddress,
+          message: 'Test message',
+          action: ComplianceAction.ONBOARD,
+          signedMessage: sha256(Buffer.from('msg')),
+          pubkey: new Uint8Array([/* public key bytes */]),
+          timestamp: 1620000000, // Valid timestamp
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.BLOCKED,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.BLOCKED);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should set status to FIRST_STRIKE for CONNECT action from a restricted country with no existing compliance status', async () => {
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          address: testConstants.defaultAddress,
+          message: 'Test message',
+          action: ComplianceAction.CONNECT,
+          signedMessage: sha256(Buffer.from('msg')),
+          pubkey: new Uint8Array([/* public key bytes */]),
+          timestamp: 1620000000, // Valid timestamp
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.FIRST_STRIKE);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should set status to COMPLIANT for any action from a non-restricted country with no existing compliance status', async () => {
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      isRestrictedCountryHeadersSpy.mockReturnValue(false);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          address: testConstants.defaultAddress,
+          message: 'Test message',
+          action: ComplianceAction.ONBOARD, // Or CONNECT, should work the same
+          signedMessage: sha256(Buffer.from('msg')),
+          pubkey: new Uint8Array([/* public key bytes */]),
+          timestamp: 1620000000, // Valid timestamp
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.COMPLIANT,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.COMPLIANT);
+    });
+
+    it('should update status to FIRST_STRIKE for CONNECT action from a restricted country with existing COMPLIANT status', async () => {
+      await ComplianceStatusTable.create({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.COMPLIANT,
+      });
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          address: testConstants.defaultAddress,
+          message: 'Test message',
+          action: ComplianceAction.CONNECT,
+          signedMessage: sha256(Buffer.from('msg')),
+          pubkey: new Uint8Array([/* public key bytes */]),
+          timestamp: 1620000000, // Valid timestamp
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.FIRST_STRIKE);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should update status to CLOSE_ONLY for CONNECT action from a restricted country with existing FIRST_STRIKE status', async () => {
+      await ComplianceStatusTable.create({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      });
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          address: testConstants.defaultAddress,
+          message: 'Test message',
+          action: ComplianceAction.CONNECT,
+          signedMessage: sha256(Buffer.from('msg')),
+          pubkey: new Uint8Array([/* public key bytes */]),
+          timestamp: 1620000000, // Valid timestamp
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.CLOSE_ONLY,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.CLOSE_ONLY);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
   });
 });