🐛 Add finalizeGenesisState to state machine

framework/src/modules/dpos_v2/module.ts

@@ -341,8 +341,15 @@ export class DPoSModule extends BaseModule {
 			},
 			genesisDataStoreSchema,
 		);
+	}
 
-		// TODO: Move below logic to finalizeGenesisState step
+	public async finalizeGenesisState(context: GenesisBlockExecuteContext): Promise<void> {
+		const assetBytes = context.assets.getAsset(this.id);
+		// if there is no asset, do not initialize
+		if (!assetBytes) {
+			return;
+		}
+		const genesisStore = codec.decode<GenesisStore>(genesisStoreSchema, assetBytes);
 		const apiContext = context.getAPIContext();
 		for (const dposValidator of genesisStore.validators) {
 			const valid = await this._validatorsAPI.registerValidatorKeys(
@@ -356,6 +363,7 @@ export class DPoSModule extends BaseModule {
 				throw new Error('Invalid validator key.');
 			}
 		}
+		const voterStore = context.getStore(this.id, STORE_PREFIX_VOTER);
 		const allVoters = await voterStore.iterateWithSchema<VoterData>(
 			{
 				start: Buffer.alloc(20),
@@ -393,6 +401,7 @@ export class DPoSModule extends BaseModule {
 		await this._validatorsAPI.setGeneratorList(apiContext, initDelegates);
 
 		const MAX_UINT32 = 2 ** 32 - 1;
+		const snapshotStore = context.getStore(this.id, STORE_PREFIX_SNAPSHOT);
 		const allSnapshots = await snapshotStore.iterate({
 			start: intToBuffer(0, 4),
 			end: intToBuffer(MAX_UINT32, 4),

framework/src/node/state_machine/state_machine.ts

@@ -41,6 +41,7 @@ export interface StateMachineModule {
 	commands: StateMachineCommand[];
 	verifyTransaction?: (ctx: TransactionVerifyContext) => Promise<VerificationResult>;
 	initGenesisState?: (ctx: GenesisBlockExecuteContext) => Promise<void>;
+	finalizeGenesisState?: (ctx: GenesisBlockExecuteContext) => Promise<void>;
 	verifyAssets?: (ctx: BlockVerifyContext) => Promise<void>;
 	beforeTransactionsExecute?: (ctx: BlockExecuteContext) => Promise<void>;
 	afterTransactionsExecute?: (ctx: BlockAfterExecuteContext) => Promise<void>;
@@ -83,6 +84,16 @@ export class StateMachine {
 				await mod.initGenesisState(blockContext);
 			}
 		}
+		for (const mod of this._modules) {
+			if (mod.finalizeGenesisState) {
+				await mod.finalizeGenesisState(blockContext);
+			}
+		}
+		for (const mod of this._systemModules) {
+			if (mod.finalizeGenesisState) {
+				await mod.finalizeGenesisState(blockContext);
+			}
+		}
 	}
 
 	public async verifyTransaction(ctx: TransactionContext): Promise<VerificationResult> {

framework/test/unit/modules/dpos_v2/module.spec.ts

@@ -127,8 +127,10 @@ describe('DPoS module', () => {
 					header: createFakeBlockHeader({ height: 12345 }),
 					assets: new BlockAssets([{ moduleID: dpos.id, data: assetBytes }]),
 				}).createGenesisBlockExecuteContext();
+				jest.spyOn(dpos, 'finalizeGenesisState');
 
 				await expect(dpos.initGenesisState(context)).rejects.toThrow(errString as string);
+				expect(dpos.finalizeGenesisState).not.toHaveBeenCalled();
 			});
 		});
 
@@ -152,7 +154,8 @@ describe('DPoS module', () => {
 					stateStore,
 					assets: new BlockAssets([{ moduleID: dpos.id, data: assetBytes }]),
 				}).createGenesisBlockExecuteContext();
-				await expect(dpos.initGenesisState(context)).rejects.toThrow(
+				await dpos.initGenesisState(context);
+				await expect(dpos.finalizeGenesisState(context)).rejects.toThrow(
 					'When genensis height is zero, there should not be a snapshot',
 				);
 			});
@@ -166,7 +169,8 @@ describe('DPoS module', () => {
 					header: createFakeBlockHeader({ height: 12345 }),
 					assets: new BlockAssets([{ moduleID: dpos.id, data: assetBytes }]),
 				}).createGenesisBlockExecuteContext();
-				await expect(dpos.initGenesisState(context)).rejects.toThrow(
+				await dpos.initGenesisState(context);
+				await expect(dpos.finalizeGenesisState(context)).rejects.toThrow(
 					'When genesis height is non-zero, snapshot is required',
 				);
 			});
@@ -185,6 +189,7 @@ describe('DPoS module', () => {
 
 			it('should store self vote and received votes', async () => {
 				await expect(dpos.initGenesisState(context)).toResolve();
+				await expect(dpos.finalizeGenesisState(context)).toResolve();
 
 				const delegateStore = stateStore.getStore(dpos.id, STORE_PREFIX_DELEGATE);
 				await expect(
@@ -255,6 +260,7 @@ describe('DPoS module', () => {
 
 			it('should register all the validators', async () => {
 				await expect(dpos.initGenesisState(context)).toResolve();
+				await expect(dpos.finalizeGenesisState(context)).toResolve();
 
 				expect(dpos['_validatorsAPI'].setGeneratorList).toHaveBeenCalledWith(
 					expect.anything(),
@@ -264,6 +270,7 @@ describe('DPoS module', () => {
 
 			it('should register all active delegates as BFT validators', async () => {
 				await expect(dpos.initGenesisState(context)).toResolve();
+				await expect(dpos.finalizeGenesisState(context)).toResolve();
 				expect(dpos['_bftAPI'].setBFTParameters).toHaveBeenCalledWith(
 					expect.anything(),
 					BigInt(68),
@@ -278,13 +285,17 @@ describe('DPoS module', () => {
 			it('should fail if registerValidatorKeys return false', async () => {
 				(dpos['_validatorsAPI'].registerValidatorKeys as jest.Mock).mockResolvedValue(false);
 
-				await expect(dpos.initGenesisState(context)).rejects.toThrow('Invalid validator key');
+				await expect(dpos.initGenesisState(context)).toResolve();
+				await expect(dpos.finalizeGenesisState(context)).rejects.toThrow('Invalid validator key');
 			});
 
 			it('should fail if getLockedAmount return different value', async () => {
 				(dpos['_tokenAPI'].getLockedAmount as jest.Mock).mockResolvedValue(BigInt(0));
 
-				await expect(dpos.initGenesisState(context)).rejects.toThrow('Voted amount is not locked');
+				await expect(dpos.initGenesisState(context)).toResolve();
+				await expect(dpos.finalizeGenesisState(context)).rejects.toThrow(
+					'Voted amount is not locked',
+				);
 			});
 		});
 	});