From ec6e56be57dff6151eee30b84a4045165566d1a0 Mon Sep 17 00:00:00 2001
From: Corey Rice <coreyarice@gmail.com>
Date: Tue, 3 Dec 2024 10:58:29 -0300
Subject: [PATCH] fix: update total delegate count, set balance to zero instead
 of deleting

---
 .../venus-governance/src/mappings/xvsVault.ts | 21 +++++++++++++------
 .../src/operations/getOrCreate.ts             | 10 +--------
 .../tests/unit/XVSVault/index.test.ts         |  6 +++---
 3 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/subgraphs/venus-governance/src/mappings/xvsVault.ts b/subgraphs/venus-governance/src/mappings/xvsVault.ts
index 3c529496..4a62a6de 100644
--- a/subgraphs/venus-governance/src/mappings/xvsVault.ts
+++ b/subgraphs/venus-governance/src/mappings/xvsVault.ts
@@ -2,7 +2,6 @@
 
 /* eslint-disable @typescript-eslint/no-empty-function */
 import { BigInt } from '@graphprotocol/graph-ts';
-import { store } from '@graphprotocol/graph-ts';
 
 import {
   DelegateChangedV2,
@@ -13,6 +12,8 @@ import {
 import { xvsVaultPid } from '../constants/config';
 import { getOrCreateDelegate } from '../operations/getOrCreate';
 import { updateDelegateChanged, updateDelegateVoteChanged } from '../operations/update';
+import { BIGINT_ONE, BIGINT_ZERO } from '../constants';
+import { getGovernanceEntity } from '../operations/get';
 
 // - event: DelegateChanged(indexed address,indexed address,indexed address)
 //   handler: handleDelegateChanged
@@ -35,8 +36,15 @@ export function handleDeposit(event: Deposit): void {
     const amount = event.params.amount;
     // Update user's staked XVS
     const result = getOrCreateDelegate(user);
+    const previousStake = result.entity.stakedXvsMantissa;
     result.entity.stakedXvsMantissa = result.entity.stakedXvsMantissa.plus(amount);
     result.entity.save();
+
+    if (previousStake.equals(BIGINT_ZERO) && amount.gt(BIGINT_ZERO)) {
+      const governance = getGovernanceEntity();
+      governance.totalDelegates = governance.totalDelegates.plus(BIGINT_ONE);
+      governance.save();
+    }
   }
 }
 
@@ -47,12 +55,13 @@ export function handleRequestedWithdrawal(event: RequestedWithdrawal): void {
     const amount = event.params.amount;
     const result = getOrCreateDelegate(user);
     const newAmount = result.entity.stakedXvsMantissa.minus(amount);
+    result.entity.stakedXvsMantissa = newAmount;
+    result.entity.save();
     // Update their delegate
-    if (newAmount.equals(new BigInt(0))) {
-      store.remove('Delegate', user.toHex());
-    } else {
-      result.entity.stakedXvsMantissa = newAmount;
-      result.entity.save();
+    if (newAmount.equals(BIGINT_ZERO)) {
+      const governance = getGovernanceEntity();
+      governance.totalDelegates = governance.totalDelegates.minus(BIGINT_ONE);
+      governance.save();
     }
   }
 }
diff --git a/subgraphs/venus-governance/src/operations/getOrCreate.ts b/subgraphs/venus-governance/src/operations/getOrCreate.ts
index f28da150..d574f46d 100644
--- a/subgraphs/venus-governance/src/operations/getOrCreate.ts
+++ b/subgraphs/venus-governance/src/operations/getOrCreate.ts
@@ -7,8 +7,7 @@ import {
   Transaction,
   TrustedRemote,
 } from '../../generated/schema';
-import { BIGINT_ONE, BIGINT_ZERO } from '../constants';
-import { nullAddress } from '../constants/addresses';
+import { BIGINT_ZERO } from '../constants';
 import { getRemoteProposalId } from '../utilities/ids';
 import {
   getDelegateId,
@@ -16,7 +15,6 @@ import {
   getRemoteProposalStateTransactionId,
   getTrustedRemoteId,
 } from '../utilities/ids';
-import { getGovernanceEntity } from './get';
 
 export class GetOrCreateDelegateReturn {
   entity: Delegate;
@@ -33,12 +31,6 @@ export const getOrCreateDelegate = (address: Address): GetOrCreateDelegateReturn
     delegate.totalVotesMantissa = BIGINT_ZERO;
     delegate.delegateCount = 0;
 
-    if (id != nullAddress) {
-      const governance = getGovernanceEntity();
-      governance.totalDelegates = governance.totalDelegates.plus(BIGINT_ONE);
-      governance.save();
-    }
-
     delegate.save();
     created = true;
   }
diff --git a/subgraphs/venus-governance/tests/unit/XVSVault/index.test.ts b/subgraphs/venus-governance/tests/unit/XVSVault/index.test.ts
index caf72b0d..9ffc986d 100644
--- a/subgraphs/venus-governance/tests/unit/XVSVault/index.test.ts
+++ b/subgraphs/venus-governance/tests/unit/XVSVault/index.test.ts
@@ -76,7 +76,7 @@ describe('XVS Vault', () => {
     assert.fieldEquals('Delegate', user.toHex(), 'stakedXvsMantissa', '800000000000000000');
   });
 
-  test('removes delegate after withdrawing everything', () => {
+  test('updates delegate after withdrawing everything', () => {
     const user = Address.fromString('0x0000000000000000000000000000000000000404');
     const amount = '800000000000000000';
     /** run handler */
@@ -89,8 +89,8 @@ describe('XVS Vault', () => {
 
     handleRequestedWithdrawal(withdrawRequestedEvent);
 
-    // Expect delegate to have been removed
-    assert.entityCount('Delegate', 1);
+    assert.entityCount('Delegate', 2);
+    assert.fieldEquals('Delegate', user.toHex(), 'stakedXvsMantissa', '0');
   });
 
   test('delegate changed', () => {