@@ -314,6 +314,58 @@ func Test_CreatePreStatesWithRegisterProtocol(t *testing.T) {
314314 require .NoError (p .CreatePreStates (ctx , sm ))
315315}
316316
317+ func TestCreatePreStatesMigration (t * testing.T ) {
318+ require := require .New (t )
319+ ctrl := gomock .NewController (t )
320+ sm := testdb .NewMockStateManager (ctrl )
321+ g := genesis .TestDefault ()
322+ mockView := NewMockContractStakeView (ctrl )
323+ mockContractStaking := NewMockContractStakingIndexer (ctrl )
324+ mockContractStaking .EXPECT ().ContractAddress ().Return (identityset .Address (1 )).Times (1 )
325+ mockContractStaking .EXPECT ().LoadStakeView (gomock .Any (), gomock .Any ()).Return (mockView , nil ).Times (1 )
326+ mockContractStaking .EXPECT ().StartHeight ().Return (uint64 (0 )).Times (1 )
327+ mockContractStaking .EXPECT ().Height ().Return (uint64 (0 ), nil ).Times (1 )
328+ p , err := NewProtocol (HelperCtx {
329+ DepositGas : nil ,
330+ BlockInterval : getBlockInterval ,
331+ }, & BuilderConfig {
332+ Staking : g .Staking ,
333+ PersistStakingPatchBlock : math .MaxUint64 ,
334+ SkipContractStakingViewHeight : math .MaxUint64 ,
335+ Revise : ReviseConfig {
336+ VoteWeight : g .Staking .VoteWeightCalConsts ,
337+ ReviseHeights : []uint64 {g .GreenlandBlockHeight }},
338+ }, nil , nil , nil , mockContractStaking )
339+ require .NoError (err )
340+ ctx := genesis .WithGenesisContext (context .Background (), g )
341+ ctx = protocol .WithBlockCtx (
342+ ctx ,
343+ protocol.BlockCtx {
344+ BlockHeight : g .ToBeEnabledBlockHeight ,
345+ },
346+ )
347+ ctx = protocol .WithFeatureCtx (protocol .WithFeatureWithHeightCtx (ctx ))
348+ v , err := p .Start (ctx , sm )
349+ require .NoError (err )
350+ require .NoError (sm .WriteView (_protocolID , v ))
351+ mockView .EXPECT ().Migrate (gomock .Any (), gomock .Any ()).Return (errors .New ("migration error" )).Times (1 )
352+ require .ErrorContains (p .CreatePreStates (ctx , sm ), "migration error" )
353+ mockView .EXPECT ().Migrate (gomock .Any (), gomock .Any ()).Return (nil ).Times (1 )
354+ require .NoError (p .CreatePreStates (ctx , sm ))
355+ require .NoError (p .CreatePreStates (protocol .WithBlockCtx (
356+ ctx ,
357+ protocol.BlockCtx {
358+ BlockHeight : g .ToBeEnabledBlockHeight - 1 ,
359+ },
360+ ), sm ))
361+ require .NoError (p .CreatePreStates (protocol .WithBlockCtx (
362+ ctx ,
363+ protocol.BlockCtx {
364+ BlockHeight : g .ToBeEnabledBlockHeight + 1 ,
365+ },
366+ ), sm ))
367+ }
368+
317369func Test_CreateGenesisStates (t * testing.T ) {
318370 require := require .New (t )
319371 ctrl := gomock .NewController (t )
0 commit comments