github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/vm/invocation_tree_test.go (about)

     1  package vm
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/nspcc-dev/neo-go/pkg/util"
     7  	"github.com/nspcc-dev/neo-go/pkg/vm/invocations"
     8  	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestInvocationTree(t *testing.T) {
    13  	script := []byte{
    14  		byte(opcode.PUSH3), byte(opcode.DEC),
    15  		byte(opcode.DUP), byte(opcode.PUSH0), byte(opcode.JMPEQ), (2 + 2 + 2 + 6 + 1),
    16  		byte(opcode.CALL), (2 + 2), // CALL shouldn't affect invocation tree.
    17  		byte(opcode.JMP), 0xf9, // DEC
    18  		byte(opcode.SYSCALL), 0, 0, 0, 0, byte(opcode.DROP),
    19  		byte(opcode.RET),
    20  		byte(opcode.RET),
    21  		byte(opcode.PUSHINT8), 0xff,
    22  	}
    23  
    24  	cnt := 0
    25  	v := newTestVM()
    26  	v.SyscallHandler = func(v *VM, _ uint32) error {
    27  		if len(v.Istack()) > 4 { // top -> call -> syscall -> call -> syscall -> ...
    28  			v.Estack().PushVal(1)
    29  			return nil
    30  		}
    31  		cnt++
    32  		v.LoadScriptWithHash(script, util.Uint160{byte(cnt)}, 0)
    33  		return nil
    34  	}
    35  	v.EnableInvocationTree()
    36  	v.LoadScript(script)
    37  	topHash := v.Context().ScriptHash()
    38  	require.NoError(t, v.Run())
    39  
    40  	res := &invocations.Tree{
    41  		Calls: []*invocations.Tree{{
    42  			Current: topHash,
    43  			Calls: []*invocations.Tree{
    44  				{
    45  					Current: util.Uint160{1},
    46  					Calls: []*invocations.Tree{
    47  						{
    48  							Current: util.Uint160{2},
    49  						},
    50  						{
    51  							Current: util.Uint160{3},
    52  						},
    53  					},
    54  				},
    55  				{
    56  					Current: util.Uint160{4},
    57  					Calls: []*invocations.Tree{
    58  						{
    59  							Current: util.Uint160{5},
    60  						},
    61  						{
    62  							Current: util.Uint160{6},
    63  						},
    64  					},
    65  				},
    66  			},
    67  		}},
    68  	}
    69  	require.Equal(t, res, v.GetInvocationTree())
    70  }