From b379b471de234e56709f01808763b4b82e26100e Mon Sep 17 00:00:00 2001
From: Travis Person <travisperson@users.noreply.github.com>
Date: Thu, 16 Feb 2023 18:57:16 +0000
Subject: [PATCH] fixup! feat: stmgr: add env to disable premigrations

---
 chain/stmgr/forks.go      |  2 +-
 chain/stmgr/forks_test.go | 66 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/chain/stmgr/forks.go b/chain/stmgr/forks.go
index 1567d036da3..5a38fbada02 100644
--- a/chain/stmgr/forks.go
+++ b/chain/stmgr/forks.go
@@ -224,7 +224,7 @@ func runPreMigration(ctx context.Context, sm *StateManager, fn PreMigrationFunc,
 	parent := ts.ParentState()
 
 	if disabled := os.Getenv(EnvDisablePreMigrations); strings.TrimSpace(disabled) == "1" {
-		log.Warn("SKIPPING pre-migration", "height", height)
+		log.Warnw("SKIPPING pre-migration", "height", height)
 		return
 	}
 
diff --git a/chain/stmgr/forks_test.go b/chain/stmgr/forks_test.go
index f56fc5ed1e7..c83c0911b36 100644
--- a/chain/stmgr/forks_test.go
+++ b/chain/stmgr/forks_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"os"
 	"sync"
 	"testing"
 
@@ -535,3 +536,68 @@ func TestForkPreMigration(t *testing.T) {
 	// to this channel.
 	require.Equal(t, 6, len(counter))
 }
+
+func TestDisablePreMigration(t *testing.T) {
+	logging.SetAllLoggers(logging.LevelInfo)
+
+	cg, err := gen.NewGenerator()
+	require.NoError(t, err)
+
+	os.Setenv(EnvDisablePreMigrations, "1")
+	defer func() {
+		os.Unsetenv(EnvDisablePreMigrations)
+	}()
+
+	counter := make(chan struct{}, 10)
+
+	sm, err := NewStateManager(
+		cg.ChainStore(),
+		consensus.NewTipSetExecutor(filcns.RewardFunc),
+		cg.StateManager().VMSys(),
+		UpgradeSchedule{{
+			Network: network.Version1,
+			Height:  testForkHeight,
+			Migration: func(_ context.Context, _ *StateManager, _ MigrationCache, _ ExecMonitor,
+				root cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) (cid.Cid, error) {
+
+				counter <- struct{}{}
+
+				return root, nil
+			},
+			PreMigrations: []PreMigration{{
+				StartWithin: 20,
+				PreMigration: func(ctx context.Context, _ *StateManager, _ MigrationCache,
+					_ cid.Cid, _ abi.ChainEpoch, _ *types.TipSet) error {
+					panic("should be skipped")
+					return nil
+				},
+			}}},
+		},
+		cg.BeaconSchedule(),
+	)
+	require.NoError(t, err)
+	require.NoError(t, sm.Start(context.Background()))
+	defer func() {
+		require.NoError(t, sm.Stop(context.Background()))
+	}()
+
+	inv := consensus.NewActorRegistry()
+	registry := builtin.MakeRegistryLegacy([]rtt.VMActor{testActor{}})
+	inv.Register(actorstypes.Version0, nil, registry)
+
+	sm.SetVMConstructor(func(ctx context.Context, vmopt *vm.VMOpts) (vm.Interface, error) {
+		nvm, err := vm.NewLegacyVM(ctx, vmopt)
+		require.NoError(t, err)
+		nvm.SetInvoker(inv)
+		return nvm, nil
+	})
+
+	cg.SetStateManager(sm)
+
+	for i := 0; i < 50; i++ {
+		_, err := cg.NextTipSet()
+		require.NoError(t, err)
+	}
+
+	require.Equal(t, 1, len(counter))
+}