pkg/morph/client/client.go

@@ -19,6 +19,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/gas"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/nep17"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/rolemgmt"
+	"github.com/nspcc-dev/neo-go/pkg/rpcclient/unwrap"
 	sc "github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/trigger"
 	"github.com/nspcc-dev/neo-go/pkg/util"
@@ -200,6 +201,45 @@ func (c *Client) Invoke(contract util.Uint160, fee fixedn.Fixed8, method string,
 	return nil
 }
 
+// TestInvokeIterator invokes contract method returning an iterator and executes cb on each element.
+// If cb returns an error, the session is closed and this error is returned as-is.
+// If the remove neo-go node does not support sessions, `unwrap.ErrNoSessionID` is returned.
+func (c *Client) TestInvokeIterator(cb func(stackitem.Item) error, contract util.Uint160, method string, args ...interface{}) error {
+	c.switchLock.RLock()
+	defer c.switchLock.RUnlock()
+
+	if c.inactive {
+		return ErrConnectionLost
+	}
+
+	val, err := c.rpcActor.Call(contract, method, args...)
+	if err != nil {
+		return err
+	} else if val.State != HaltState {
+		return wrapFrostFSError(&notHaltStateError{state: val.State, exception: val.FaultException})
+	}
+
+	sid, r, err := unwrap.SessionIterator(val, err)
+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		_ = c.rpcActor.TerminateSession(sid)
+	}()
+
+	items, err := c.rpcActor.TraverseIterator(sid, &r, 0)
+	for err == nil && len(items) != 0 {
+		for i := range items {
+			if err = cb(items[i]); err != nil {
+				return err
+			}
+		}
+		items, err = c.rpcActor.TraverseIterator(sid, &r, 0)
+	}
+	return err
+}
+
 // TestInvoke invokes contract method locally in neo-go node. This method should
 // be used to read data from smart-contract.
 func (c *Client) TestInvoke(contract util.Uint160, method string, args ...any) (res []stackitem.Item, err error) {