Merge pull request #2298 from nspcc-dev/test-signer

Improve signers handling in test framework

pkg/neotest/chain/chain.go

@@ -72,40 +72,36 @@ func init() {
 	standByCommittee[4] = hex.EncodeToString(pubs[4].Bytes())
 	standByCommittee[5] = hex.EncodeToString(pubs[5].Bytes())
 
-	multiValidatorAcc = make([]*wallet.Account, mv)
+	multiValidatorAcc = make([]*wallet.Account, 4)
 	sort.Sort(pubs[:4])
 
-vloop:
-	for i := 0; i < mv; i++ {
-		for j := range accs {
-			if accs[j].PrivateKey().PublicKey().Equal(pubs[i]) {
-				multiValidatorAcc[i] = wallet.NewAccountFromPrivateKey(accs[j].PrivateKey())
-				err := multiValidatorAcc[i].ConvertMultisig(mv, pubs[:4])
-				if err != nil {
-					panic(err)
-				}
-				continue vloop
-			}
+	sort.Slice(accs[:4], func(i, j int) bool {
+		p1 := accs[i].PrivateKey().PublicKey()
+		p2 := accs[j].PrivateKey().PublicKey()
+		return p1.Cmp(p2) == -1
+	})
+	for i := range multiValidatorAcc {
+		multiValidatorAcc[i] = wallet.NewAccountFromPrivateKey(accs[i].PrivateKey())
+		err := multiValidatorAcc[i].ConvertMultisig(mv, pubs[:4])
+		if err != nil {
+			panic(err)
 		}
-		panic("invalid committee WIFs")
 	}
 
-	multiCommitteeAcc = make([]*wallet.Account, mc)
+	multiCommitteeAcc = make([]*wallet.Account, len(committeeWIFs))
 	sort.Sort(pubs)
 
-cloop:
-	for i := 0; i < mc; i++ {
-		for j := range accs {
-			if accs[j].PrivateKey().PublicKey().Equal(pubs[i]) {
-				multiCommitteeAcc[i] = wallet.NewAccountFromPrivateKey(accs[j].PrivateKey())
-				err := multiCommitteeAcc[i].ConvertMultisig(mc, pubs)
-				if err != nil {
-					panic(err)
-				}
-				continue cloop
-			}
+	sort.Slice(accs, func(i, j int) bool {
+		p1 := accs[i].PrivateKey().PublicKey()
+		p2 := accs[j].PrivateKey().PublicKey()
+		return p1.Cmp(p2) == -1
+	})
+	for i := range multiCommitteeAcc {
+		multiCommitteeAcc[i] = wallet.NewAccountFromPrivateKey(accs[i].PrivateKey())
+		err := multiCommitteeAcc[i].ConvertMultisig(mc, pubs)
+		if err != nil {
+			panic(err)
 		}
-		panic("invalid committee WIFs")
 	}
 }
 

pkg/neotest/signer.go

@@ -3,6 +3,7 @@ package neotest
 import (
 	"bytes"
 	"fmt"
+	"sort"
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/pkg/config/netmode"
@@ -27,15 +28,33 @@ type Signer interface {
 	SignTx(netmode.Magic, *transaction.Transaction) error
 }
 
+// SingleSigner is a generic interface for simple one-signature signer.
+type SingleSigner interface {
+	Signer
+	// Account returns underlying account which can be used to
+	// get public key and/or sign arbitrary things.
+	Account() *wallet.Account
+}
+
+// MultiSigner is the interface for multisignature signing account.
+type MultiSigner interface {
+	Signer
+	// Single returns simple-signature signer for n-th account in list.
+	Single(n int) SingleSigner
+}
+
 // signer represents simple-signature signer.
 type signer wallet.Account
 
 // multiSigner represents single multi-signature signer consisting of provided accounts.
-type multiSigner []*wallet.Account
+type multiSigner struct {
+	accounts []*wallet.Account
+	m        int
+}
 
 // NewSingleSigner returns multi-signature signer for the provided account.
 // It must contain exactly as many accounts as needed to sign the script.
-func NewSingleSigner(acc *wallet.Account) Signer {
+func NewSingleSigner(acc *wallet.Account) SingleSigner {
 	if !vm.IsSignatureContract(acc.Contract.Script) {
 		panic("account must have simple-signature verification script")
 	}
@@ -63,9 +82,14 @@ func (s *signer) SignTx(magic netmode.Magic, tx *transaction.Transaction) error
 	return (*wallet.Account)(s).SignTx(magic, tx)
 }
 
+// Account implements SingleSigner interface.
+func (s *signer) Account() *wallet.Account {
+	return (*wallet.Account)(s)
+}
+
 // NewMultiSigner returns multi-signature signer for the provided account.
 // It must contain at least as many accounts as needed to sign the script.
-func NewMultiSigner(accs ...*wallet.Account) Signer {
+func NewMultiSigner(accs ...*wallet.Account) MultiSigner {
 	if len(accs) == 0 {
 		panic("empty account list")
 	}
@@ -78,30 +102,35 @@ func NewMultiSigner(accs ...*wallet.Account) Signer {
 		panic(fmt.Sprintf("verification script requires %d signatures, "+
 			"but only %d accounts were provided", m, len(accs)))
 	}
+	sort.Slice(accs, func(i, j int) bool {
+		p1 := accs[i].PrivateKey().PublicKey()
+		p2 := accs[j].PrivateKey().PublicKey()
+		return p1.Cmp(p2) == -1
+	})
 	for _, acc := range accs {
 		if !bytes.Equal(script, acc.Contract.Script) {
 			panic("all accounts must have equal verification script")
 		}
 	}
 
-	return multiSigner(accs[:m])
+	return multiSigner{accounts: accs, m: m}
 }
 
 // ScriptHash implements Signer interface.
 func (m multiSigner) ScriptHash() util.Uint160 {
-	return m[0].Contract.ScriptHash()
+	return m.accounts[0].Contract.ScriptHash()
 }
 
 // Script implements Signer interface.
 func (m multiSigner) Script() []byte {
-	return m[0].Contract.Script
+	return m.accounts[0].Contract.Script
 }
 
 // SignHashable implements Signer interface.
 func (m multiSigner) SignHashable(magic uint32, item hash.Hashable) []byte {
 	var script []byte
-	for _, acc := range m {
-		sign := acc.PrivateKey().SignHashable(magic, item)
+	for i := 0; i < m.m; i++ {
+		sign := m.accounts[i].PrivateKey().SignHashable(magic, item)
 		script = append(script, byte(opcode.PUSHDATA1), 64)
 		script = append(script, sign...)
 	}
@@ -125,9 +154,19 @@ func (m multiSigner) SignTx(magic netmode.Magic, tx *transaction.Transaction) er
 	return nil
 }
 
+// Single implements MultiSigner interface.
+func (m multiSigner) Single(n int) SingleSigner {
+	if len(m.accounts) <= n {
+		panic("invalid index")
+	}
+	return NewSingleSigner(wallet.NewAccountFromPrivateKey(m.accounts[n].PrivateKey()))
+}
+
 func checkMultiSigner(t *testing.T, s Signer) {
-	accs, ok := s.(multiSigner)
+	ms, ok := s.(multiSigner)
 	require.True(t, ok, "expected to be a multi-signer")
+
+	accs := ms.accounts
 	require.True(t, len(accs) > 0, "empty multi-signer")
 
 	m := len(accs[0].Contract.Parameters)

pkg/neotest/signer_test.go

@@ -0,0 +1,48 @@
+package neotest
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
+	"github.com/nspcc-dev/neo-go/pkg/wallet"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSingleSigner(t *testing.T) {
+	a, err := wallet.NewAccount()
+	require.NoError(t, err)
+
+	s := NewSingleSigner(a)
+	require.Equal(t, s.ScriptHash(), s.Account().Contract.ScriptHash())
+}
+
+func TestMultiSigner(t *testing.T) {
+	const size = 4
+
+	pubs := make(keys.PublicKeys, size)
+	accs := make([]*wallet.Account, size)
+	for i := range accs {
+		a, err := wallet.NewAccount()
+		require.NoError(t, err)
+
+		accs[i] = a
+		pubs[i] = a.PrivateKey().PublicKey()
+	}
+
+	sort.Sort(pubs)
+	m := smartcontract.GetDefaultHonestNodeCount(size)
+	for i := range accs {
+		require.NoError(t, accs[i].ConvertMultisig(m, pubs))
+	}
+
+	s := NewMultiSigner(accs...)
+	for i := range pubs {
+		for j := range accs {
+			if pub := accs[j].PrivateKey().PublicKey(); pub.Equal(pubs[i]) {
+				require.Equal(t, pub, s.Single(i).Account().PrivateKey().PublicKey())
+			}
+		}
+	}
+}