package tutustest import ( "math/big" "testing" "time" "github.com/stretchr/testify/require" "git.marketally.com/tutus-one/tutus-chain/pkg/core/native/nativenames" "git.marketally.com/tutus-one/tutus-chain/pkg/core/transaction" "git.marketally.com/tutus-one/tutus-chain/pkg/util" "git.marketally.com/tutus-one/tutus-chain/pkg/vm/stackitem" ) // GovernmentHelper provides utilities for testing Tutus government contracts. // It wraps common operations like registering citizens, granting roles, and // setting up cross-contract test scenarios. type GovernmentHelper struct { *Executor t testing.TB // Contract invokers for common government contracts Vita *ContractInvoker Lex *ContractInvoker Eligere *ContractInvoker Scire *ContractInvoker Salus *ContractInvoker VTS *ContractInvoker Annos *ContractInvoker Tribute *ContractInvoker RoleReg *ContractInvoker Treasury *ContractInvoker } // NewGovernmentHelper creates a helper for testing government contracts. // It initializes invokers for all major government contracts. func NewGovernmentHelper(t testing.TB, e *Executor) *GovernmentHelper { g := &GovernmentHelper{ Executor: e, t: t, } // Initialize contract invokers with committee authority g.Vita = e.CommitteeInvoker(e.NativeHash(t, nativenames.Vita)) g.Lex = e.CommitteeInvoker(e.NativeHash(t, nativenames.Lex)) g.Eligere = e.CommitteeInvoker(e.NativeHash(t, nativenames.Eligere)) g.Scire = e.CommitteeInvoker(e.NativeHash(t, nativenames.Scire)) g.Salus = e.CommitteeInvoker(e.NativeHash(t, nativenames.Salus)) g.VTS = e.CommitteeInvoker(e.NativeHash(t, nativenames.VTS)) g.Annos = e.CommitteeInvoker(e.NativeHash(t, nativenames.Annos)) g.Tribute = e.CommitteeInvoker(e.NativeHash(t, nativenames.Tribute)) g.RoleReg = e.CommitteeInvoker(e.NativeHash(t, nativenames.RoleRegistry)) g.Treasury = e.CommitteeInvoker(e.NativeHash(t, nativenames.Treasury)) return g } // Citizen represents a registered Vita holder for testing. type Citizen struct { Account SingleSigner VitaID uint64 BirthTime uint64 Registered bool } // RegisterCitizen registers a new Vita token for the given account. // birthTimestamp is the citizen's birth date (Unix timestamp in milliseconds). // Returns the Citizen with VitaID populated. func (g *GovernmentHelper) RegisterCitizen(account SingleSigner, birthTimestamp uint64) *Citizen { citizen := &Citizen{ Account: account, BirthTime: birthTimestamp, } // Get current token count to predict the new VitaID g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) { count, err := stack[0].TryInteger() require.NoError(t, err) citizen.VitaID = count.Uint64() }, "totalSupply") // Register the Vita token g.Vita.Invoke(g.t, true, "register", account.ScriptHash(), birthTimestamp) citizen.Registered = true return citizen } // RegisterAdultCitizen registers a citizen who is 25 years old (adult). // Useful for tests requiring voting age or adult status. func (g *GovernmentHelper) RegisterAdultCitizen(account SingleSigner) *Citizen { // 25 years ago in milliseconds birthTime := uint64(time.Now().AddDate(-25, 0, 0).UnixMilli()) return g.RegisterCitizen(account, birthTime) } // RegisterChildCitizen registers a citizen who is 10 years old (child). // Useful for tests verifying age restrictions. func (g *GovernmentHelper) RegisterChildCitizen(account SingleSigner) *Citizen { // 10 years ago in milliseconds birthTime := uint64(time.Now().AddDate(-10, 0, 0).UnixMilli()) return g.RegisterCitizen(account, birthTime) } // RegisterElderCitizen registers a citizen who is 70 years old (elder). // Useful for tests involving retirement age. func (g *GovernmentHelper) RegisterElderCitizen(account SingleSigner) *Citizen { // 70 years ago in milliseconds birthTime := uint64(time.Now().AddDate(-70, 0, 0).UnixMilli()) return g.RegisterCitizen(account, birthTime) } // NewCitizenAccount creates a new account and registers it as a citizen. // Returns both the citizen and their account for signing transactions. func (g *GovernmentHelper) NewCitizenAccount() *Citizen { acc := g.Vita.NewAccount(g.t).(SingleSigner) return g.RegisterAdultCitizen(acc) } // VerifyVitaOwnership checks that the given account owns a Vita token. func (g *GovernmentHelper) VerifyVitaOwnership(account util.Uint160) bool { var hasVita bool g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) { balance, err := stack[0].TryInteger() require.NoError(t, err) hasVita = balance.Cmp(big.NewInt(0)) > 0 }, "balanceOf", account) return hasVita } // GetVitaID returns the Vita token ID for the given owner, or -1 if not found. func (g *GovernmentHelper) GetVitaID(owner util.Uint160) int64 { var vitaID int64 = -1 g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) { if stack[0].Type() != stackitem.AnyT { tokens := stack[0].Value().([]stackitem.Item) if len(tokens) > 0 { id, err := tokens[0].TryInteger() require.NoError(t, err) vitaID = id.Int64() } } }, "tokensOf", owner) return vitaID } // SuspendVita suspends a citizen's Vita token (requires committee + Lex restriction). func (g *GovernmentHelper) SuspendVita(vitaID uint64, reason string) { // First, create a Lex liberty restriction (required for due process) // This would normally require a judicial order g.Vita.Invoke(g.t, true, "suspend", vitaID, reason) } // TransferVTS transfers VTS tokens between accounts. func (g *GovernmentHelper) TransferVTS(from, to SingleSigner, amount int64) util.Uint256 { vtsInvoker := g.VTS.WithSigners(from) return vtsInvoker.Invoke(g.t, true, "transfer", from.ScriptHash(), to.ScriptHash(), amount, nil) } // GetVTSBalance returns the VTS balance for an account. func (g *GovernmentHelper) GetVTSBalance(account util.Uint160) *big.Int { var balance *big.Int g.VTS.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) { var err error balance, err = stack[0].TryInteger() require.NoError(t, err) }, "balanceOf", account) return balance } // MintVTS mints VTS tokens to an account (committee only). func (g *GovernmentHelper) MintVTS(to util.Uint160, amount int64) { g.VTS.Invoke(g.t, true, "mint", to, amount) } // CreateProposal creates a new Eligere proposal. // Returns the proposal ID. func (g *GovernmentHelper) CreateProposal(proposer *Citizen, title, description string, category int64) uint64 { eligereInvoker := g.Eligere.WithSigners(proposer.Account) var proposalID uint64 eligereInvoker.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) { id, err := stack[0].TryInteger() require.NoError(t, err) proposalID = id.Uint64() }, "createProposal", title, description, category) return proposalID } // Vote casts a vote on a proposal. func (g *GovernmentHelper) Vote(voter *Citizen, proposalID uint64, vote int64) { eligereInvoker := g.Eligere.WithSigners(voter.Account) eligereInvoker.Invoke(g.t, true, "vote", proposalID, vote) } // PrepareVitaRegistration prepares a Vita registration transaction without executing. // Useful for batching or testing transaction ordering. func (g *GovernmentHelper) PrepareVitaRegistration(account SingleSigner, birthTimestamp uint64) *transaction.Transaction { return g.Vita.PrepareInvoke(g.t, "register", account.ScriptHash(), birthTimestamp) } // BatchRegisterCitizens registers multiple citizens in a single block. func (g *GovernmentHelper) BatchRegisterCitizens(accounts []SingleSigner) []*Citizen { citizens := make([]*Citizen, len(accounts)) txs := make([]*transaction.Transaction, len(accounts)) birthTime := uint64(time.Now().AddDate(-25, 0, 0).UnixMilli()) // Get starting VitaID var startID uint64 g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) { count, err := stack[0].TryInteger() require.NoError(t, err) startID = count.Uint64() }, "totalSupply") // Prepare all transactions for i, acc := range accounts { citizens[i] = &Citizen{ Account: acc, BirthTime: birthTime, VitaID: startID + uint64(i), } txs[i] = g.Vita.PrepareInvoke(g.t, "register", acc.ScriptHash(), birthTime) } // Execute all in one block g.Vita.AddNewBlock(g.t, txs...) // Verify all succeeded for i, tx := range txs { g.Vita.CheckHalt(g.t, tx.Hash(), stackitem.Make(true)) citizens[i].Registered = true } return citizens }