From 342134f8535cea4f9f2978463520e1ce48b2ed67 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Kripalani?= <raul@protocol.ai>
Date: Fri, 20 Aug 2021 18:07:25 +0100
Subject: [PATCH] EnableLargeSectors: restore previous supported proof types.

---
 itests/deals_512mb_test.go |  2 +-
 itests/kit/sectors.go      | 11 +++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/itests/deals_512mb_test.go b/itests/deals_512mb_test.go
index eac598ab9..766d83835 100644
--- a/itests/deals_512mb_test.go
+++ b/itests/deals_512mb_test.go
@@ -15,7 +15,7 @@ func TestStorageDealMissingBlock(t *testing.T) {
 	ctx := context.Background()
 
 	// enable 512MiB proofs so we can conduct larger transfers.
-	kit.EnableLargeSectors()
+	kit.EnableLargeSectors(t)
 	kit.QuietMiningLogs()
 
 	client, miner, ens := kit.EnsembleMinimal(t,
diff --git a/itests/kit/sectors.go b/itests/kit/sectors.go
index 5f82055c0..d9c52b9ca 100644
--- a/itests/kit/sectors.go
+++ b/itests/kit/sectors.go
@@ -1,6 +1,8 @@
 package kit
 
 import (
+	"testing"
+
 	"github.com/filecoin-project/go-state-types/abi"
 
 	"github.com/filecoin-project/lotus/chain/actors/policy"
@@ -8,9 +10,14 @@ import (
 
 // EnableLargeSectors enables 512MiB sectors. This is useful in combination with
 // mock proofs, for testing larger transfers.
-func EnableLargeSectors() {
+func EnableLargeSectors(t *testing.T) {
 	policy.SetSupportedProofTypes(
 		abi.RegisteredSealProof_StackedDrg2KiBV1,
-		abi.RegisteredSealProof_StackedDrg512MiBV1, // <==
+		abi.RegisteredSealProof_StackedDrg512MiBV1, // <== here
 	)
+	t.Cleanup(func() { // reset when done.
+		policy.SetSupportedProofTypes(
+			abi.RegisteredSealProof_StackedDrg2KiBV1,
+		)
+	})
 }