Add function for testing framework's blockchain to create/load snapshots

runtime/stdlib/contracts/test.cdc

@@ -134,6 +134,27 @@ pub contract Test {
         pub fun moveTime(by delta: Fix64) {
             self.backend.moveTime(by: delta)
         }
+
+        /// Creates a snapshot of the blockchain, at the
+        /// current ledger state, with the given name.
+        ///
+        pub fun createSnapshot(_ name: String) {
+            let err = self.backend.createSnapshot(name)
+            if err != nil {
+                panic(err!.message)
+            }
+        }
+
+        /// Loads a snapshot of the blockchain, with the
+        /// given name, and updates the current ledger
+        /// state.
+        ///
+        pub fun loadSnapshot(_ name: String) {
+            let err = self.backend.loadSnapshot(name)
+            if err != nil {
+                panic(err!.message)
+            }
+        }
     }
 
     pub struct Matcher {
@@ -324,6 +345,17 @@ pub contract Test {
         /// which should be passed in the form of seconds.
         ///
         pub fun moveTime(by delta: Fix64)
+
+        /// Creates a snapshot of the blockchain, at the
+        /// current ledger state, with the given name.
+        ///
+        pub fun createSnapshot(_ name: String): Error?
+
+        /// Loads a snapshot of the blockchain, with the
+        /// given name, and updates the current ledger
+        /// state.
+        ///
+        pub fun loadSnapshot(_ name: String): Error?
     }
 
     /// Returns a new matcher that negates the test of the given matcher.

runtime/stdlib/test-framework.go

@@ -79,6 +79,10 @@ type Blockchain interface {
 	Reset(uint64)
 
 	MoveTime(int64)
+
+	CreateSnapshot(string) error
+
+	LoadSnapshot(string) error
 }
 
 type ScriptResult struct {

runtime/stdlib/test_emulatorbackend.go

@@ -47,6 +47,8 @@ type testEmulatorBackendType struct {
 	eventsFunctionType                 *sema.FunctionType
 	resetFunctionType                  *sema.FunctionType
 	moveTimeFunctionType               *sema.FunctionType
+	createSnapshotFunctionType         *sema.FunctionType
+	loadSnapshotFunctionType           *sema.FunctionType
 }
 
 func newTestEmulatorBackendType(
@@ -112,6 +114,16 @@ func newTestEmulatorBackendType(
 		testEmulatorBackendTypeMoveTimeFunctionName,
 	)
 
+	createSnapshotFunctionType := interfaceFunctionType(
+		blockchainBackendInterfaceType,
+		testEmulatorBackendTypeCreateSnapshotFunctionName,
+	)
+
+	loadSnapshotFunctionType := interfaceFunctionType(
+		blockchainBackendInterfaceType,
+		testEmulatorBackendTypeLoadSnapshotFunctionName,
+	)
+
 	compositeType := &sema.CompositeType{
 		Identifier: testEmulatorBackendTypeName,
 		Kind:       common.CompositeKindStructure,
@@ -194,6 +206,18 @@ func newTestEmulatorBackendType(
 			moveTimeFunctionType,
 			testEmulatorBackendTypeMoveTimeFunctionDocString,
 		),
+		sema.NewUnmeteredPublicFunctionMember(
+			compositeType,
+			testEmulatorBackendTypeCreateSnapshotFunctionName,
+			createSnapshotFunctionType,
+			testEmulatorBackendTypeCreateSnapshotFunctionDocString,
+		),
+		sema.NewUnmeteredPublicFunctionMember(
+			compositeType,
+			testEmulatorBackendTypeLoadSnapshotFunctionName,
+			loadSnapshotFunctionType,
+			testEmulatorBackendTypeLoadSnapshotFunctionDocString,
+		),
 	}
 
 	compositeType.Members = sema.MembersAsMap(members)
@@ -213,6 +237,8 @@ func newTestEmulatorBackendType(
 		eventsFunctionType:                 eventsFunctionType,
 		resetFunctionType:                  resetFunctionType,
 		moveTimeFunctionType:               moveTimeFunctionType,
+		createSnapshotFunctionType:         createSnapshotFunctionType,
+		loadSnapshotFunctionType:           loadSnapshotFunctionType,
 	}
 }
 
@@ -739,6 +765,58 @@ func (t *testEmulatorBackendType) newMoveTimeFunction(
 	)
 }
 
+// 'Emulator.createSnapshot' function
+
+const testEmulatorBackendTypeCreateSnapshotFunctionName = "createSnapshot"
+
+const testEmulatorBackendTypeCreateSnapshotFunctionDocString = `
+Creates a snapshot of the blockchain, at the
+current ledger state, with the given name.
+`
+
+func (t *testEmulatorBackendType) newCreateSnapshotFunction(
+	blockchain Blockchain,
+) *interpreter.HostFunctionValue {
+	return interpreter.NewUnmeteredHostFunctionValue(
+		t.createSnapshotFunctionType,
+		func(invocation interpreter.Invocation) interpreter.Value {
+			name, ok := invocation.Arguments[0].(*interpreter.StringValue)
+			if !ok {
+				panic(errors.NewUnreachableError())
+			}
+
+			err := blockchain.CreateSnapshot(name.Str)
+			return newErrorValue(invocation.Interpreter, err)
+		},
+	)
+}
+
+// 'Emulator.loadSnapshot' function
+
+const testEmulatorBackendTypeLoadSnapshotFunctionName = "loadSnapshot"
+
+const testEmulatorBackendTypeLoadSnapshotFunctionDocString = `
+Loads a snapshot of the blockchain, with the given name, and
+updates the current ledger state.
+`
+
+func (t *testEmulatorBackendType) newLoadSnapshotFunction(
+	blockchain Blockchain,
+) *interpreter.HostFunctionValue {
+	return interpreter.NewUnmeteredHostFunctionValue(
+		t.loadSnapshotFunctionType,
+		func(invocation interpreter.Invocation) interpreter.Value {
+			name, ok := invocation.Arguments[0].(*interpreter.StringValue)
+			if !ok {
+				panic(errors.NewUnreachableError())
+			}
+
+			err := blockchain.LoadSnapshot(name.Str)
+			return newErrorValue(invocation.Interpreter, err)
+		},
+	)
+}
+
 func (t *testEmulatorBackendType) newEmulatorBackend(
 	inter *interpreter.Interpreter,
 	blockchain Blockchain,
@@ -792,6 +870,14 @@ func (t *testEmulatorBackendType) newEmulatorBackend(
 			Name:  testEmulatorBackendTypeMoveTimeFunctionName,
 			Value: t.newMoveTimeFunction(blockchain),
 		},
+		{
+			Name:  testEmulatorBackendTypeCreateSnapshotFunctionName,
+			Value: t.newCreateSnapshotFunction(blockchain),
+		},
+		{
+			Name:  testEmulatorBackendTypeLoadSnapshotFunctionName,
+			Value: t.newLoadSnapshotFunction(blockchain),
+		},
 	}
 
 	// TODO: Use SimpleCompositeValue

runtime/stdlib/test_test.go

@@ -20,6 +20,7 @@ package stdlib
 
 import (
 	"errors"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -2339,6 +2340,162 @@ func TestBlockchain(t *testing.T) {
 		assert.True(t, newEmulatorBackendInvoked)
 	})
 
+	t.Run("createSnapshot", func(t *testing.T) {
+		t.Parallel()
+
+		const script = `
+            import Test
+
+            pub fun test() {
+                let blockchain = Test.newEmulatorBlockchain()
+                blockchain.createSnapshot("adminCreated")
+            }
+		`
+
+		createSnapshotInvoked := false
+
+		testFramework := &mockedTestFramework{
+			newEmulatorBackend: func() Blockchain {
+				return &mockedBlockchain{
+					createSnapshot: func(name string) error {
+						createSnapshotInvoked = true
+						assert.Equal(t, "adminCreated", name)
+
+						return nil
+					},
+				}
+			},
+		}
+
+		inter, err := newTestContractInterpreterWithTestFramework(t, script, testFramework)
+		require.NoError(t, err)
+
+		_, err = inter.Invoke("test")
+		require.NoError(t, err)
+
+		assert.True(t, createSnapshotInvoked)
+	})
+
+	t.Run("createSnapshot failure", func(t *testing.T) {
+		t.Parallel()
+
+		const script = `
+            import Test
+
+            pub fun test() {
+                let blockchain = Test.newEmulatorBlockchain()
+                blockchain.createSnapshot("adminCreated")
+            }
+		`
+
+		createSnapshotInvoked := false
+
+		testFramework := &mockedTestFramework{
+			newEmulatorBackend: func() Blockchain {
+				return &mockedBlockchain{
+					createSnapshot: func(name string) error {
+						createSnapshotInvoked = true
+						assert.Equal(t, "adminCreated", name)
+
+						return fmt.Errorf("failed to create snapshot: %s", name)
+					},
+				}
+			},
+		}
+
+		inter, err := newTestContractInterpreterWithTestFramework(t, script, testFramework)
+		require.NoError(t, err)
+
+		_, err = inter.Invoke("test")
+		require.ErrorContains(t, err, "panic: failed to create snapshot: adminCreated")
+
+		assert.True(t, createSnapshotInvoked)
+	})
+
+	t.Run("loadSnapshot", func(t *testing.T) {
+		t.Parallel()
+
+		const script = `
+            import Test
+
+            pub fun test() {
+                let blockchain = Test.newEmulatorBlockchain()
+                blockchain.createSnapshot("adminCreated")
+				blockchain.loadSnapshot("adminCreated")
+            }
+		`
+
+		loadSnapshotInvoked := false
+
+		testFramework := &mockedTestFramework{
+			newEmulatorBackend: func() Blockchain {
+				return &mockedBlockchain{
+					createSnapshot: func(name string) error {
+						assert.Equal(t, "adminCreated", name)
+
+						return nil
+					},
+					loadSnapshot: func(name string) error {
+						loadSnapshotInvoked = true
+						assert.Equal(t, "adminCreated", name)
+
+						return nil
+					},
+				}
+			},
+		}
+
+		inter, err := newTestContractInterpreterWithTestFramework(t, script, testFramework)
+		require.NoError(t, err)
+
+		_, err = inter.Invoke("test")
+		require.NoError(t, err)
+
+		assert.True(t, loadSnapshotInvoked)
+	})
+
+	t.Run("loadSnapshot failure", func(t *testing.T) {
+		t.Parallel()
+
+		const script = `
+            import Test
+
+            pub fun test() {
+                let blockchain = Test.newEmulatorBlockchain()
+                blockchain.createSnapshot("adminCreated")
+				blockchain.loadSnapshot("contractDeployed")
+            }
+		`
+
+		loadSnapshotInvoked := false
+
+		testFramework := &mockedTestFramework{
+			newEmulatorBackend: func() Blockchain {
+				return &mockedBlockchain{
+					createSnapshot: func(name string) error {
+						assert.Equal(t, "adminCreated", name)
+
+						return nil
+					},
+					loadSnapshot: func(name string) error {
+						loadSnapshotInvoked = true
+						assert.Equal(t, "contractDeployed", name)
+
+						return fmt.Errorf("failed to create snapshot: %s", name)
+					},
+				}
+			},
+		}
+
+		inter, err := newTestContractInterpreterWithTestFramework(t, script, testFramework)
+		require.NoError(t, err)
+
+		_, err = inter.Invoke("test")
+		require.ErrorContains(t, err, "panic: failed to create snapshot: contractDeployed")
+
+		assert.True(t, loadSnapshotInvoked)
+	})
+
 	// TODO: Add more tests for the remaining functions.
 }
 
@@ -2379,6 +2536,8 @@ type mockedBlockchain struct {
 	events             func(inter *interpreter.Interpreter, eventType interpreter.StaticType) interpreter.Value
 	reset              func(uint64)
 	moveTime           func(int64)
+	createSnapshot     func(string) error
+	loadSnapshot       func(string) error
 }
 
 var _ Blockchain = &mockedBlockchain{}
@@ -2505,3 +2664,19 @@ func (m mockedBlockchain) MoveTime(timeDelta int64) {
 
 	m.moveTime(timeDelta)
 }
+
+func (m mockedBlockchain) CreateSnapshot(name string) error {
+	if m.createSnapshot == nil {
+		panic("'CreateSnapshot' is not implemented")
+	}
+
+	return m.createSnapshot(name)
+}
+
+func (m mockedBlockchain) LoadSnapshot(name string) error {
+	if m.loadSnapshot == nil {
+		panic("'LoadSnapshot' is not implemented")
+	}
+
+	return m.loadSnapshot(name)
+}