add functions

types/upgrade.go

@@ -45,13 +45,31 @@ func (mgr *UpgradeManager) GetUpgradeHeight(name string) int64 {
 	if mgr.Config.HeightMap == nil {
 		return 0
 	}
-	return mgr.Config.HeightMap[UpgradeLimitAddressLength]
+	return mgr.Config.HeightMap[name]
+}
+
+func IsUpgradeHeight(name string) bool {
+	upgradeHeight := UpgradeMgr.GetUpgradeHeight(name)
+	if upgradeHeight == 0 {
+		return false
+	}
+
+	return upgradeHeight == UpgradeMgr.GetHeight()
+}
+
+func IsUpgrade(name string) bool {
+	upgradeHeight := UpgradeMgr.GetUpgradeHeight(name)
+	if upgradeHeight == 0 {
+		return false
+	}
+
+	return UpgradeMgr.GetHeight() >= upgradeHeight
 }
 
 func IsLimitAddressLengthUpgrade() bool {
-	forkHeight := UpgradeMgr.GetUpgradeHeight(UpgradeLimitAddressLength)
-	if forkHeight == 0 {
+	upgradeHeight := UpgradeMgr.GetUpgradeHeight(UpgradeLimitAddressLength)
+	if upgradeHeight == 0 {
 		return false
 	}
-	return UpgradeMgr.GetHeight() >= forkHeight
+	return UpgradeMgr.GetHeight() >= upgradeHeight
 }

types/upgrade_test.go

@@ -15,48 +15,55 @@ func TestIsLimitAddressLengthFork(t *testing.T) {
 	UpgradeMgr = NewUpgradeManager(config)
 
 	type testCase struct {
-		config UpgradeConfig
-		height int64
-		result bool
+		config        UpgradeConfig
+		height        int64
+		upgradeResult bool
+		heightResult  bool
 	}
 
 	testCases := []testCase{
 		{
 			config: UpgradeConfig{
 				map[string]int64{},
 			},
-			height: 10000,
-			result: false,
+			height:        10000,
+			upgradeResult: false,
+			heightResult:  false,
 		},
 		{
 			config: UpgradeConfig{
 				map[string]int64{
 					UpgradeLimitAddressLength: 545000,
 				},
 			},
-			height: 10000,
-			result: false,
+			height:        10000,
+			upgradeResult: false,
+			heightResult:  false,
 		}, {
 			config: UpgradeConfig{
 				map[string]int64{
 					UpgradeLimitAddressLength: 545000,
 				},
 			},
-			height: 545000,
-			result: true,
+			height:        545000,
+			upgradeResult: true,
+			heightResult:  true,
 		}, {
 			config: UpgradeConfig{
 				map[string]int64{
 					UpgradeLimitAddressLength: 545000,
 				},
 			},
-			height: 545001,
-			result: true,
+			height:        545001,
+			upgradeResult: true,
+			heightResult:  false,
 		},
 	}
 
 	for _, tc := range testCases {
 		UpgradeMgr.SetHeight(tc.height)
-		require.Equal(t, tc.result, IsLimitAddressLengthUpgrade())
+		require.Equal(t, tc.upgradeResult, IsLimitAddressLengthUpgrade())
+		require.Equal(t, tc.upgradeResult, IsUpgrade(UpgradeLimitAddressLength))
+		require.Equal(t, tc.heightResult, IsUpgradeHeight(UpgradeLimitAddressLength))
 	}
 }