Refactor implicit account creation in dockertestframework

pkg/testsuite/mock/wallet.go

@@ -62,6 +62,26 @@ func NewAccountData(accountID iotago.AccountID, optAddressIndex ...uint32) *Acco
 	}
 }
 
+type AccountWithWallet struct {
+	account *AccountData
+	wallet  *Wallet
+}
+
+func NewAccountWithWallet(account *AccountData, wallet *Wallet) *AccountWithWallet {
+	return &AccountWithWallet{
+		account: account,
+		wallet:  wallet,
+	}
+}
+
+func (a *AccountWithWallet) Account() *AccountData {
+	return a.account
+}
+
+func (a *AccountWithWallet) Wallet() *Wallet {
+	return a.wallet
+}
+
 // WalletClock is an interface that provides the current slot.
 type WalletClock interface {
 	SetCurrentSlot(slot iotago.SlotIndex)

tools/docker-network/tests/api_core_test.go

@@ -282,39 +282,14 @@ func Test_ValidatorsAPI(t *testing.T) {
 	})
 
 	// Create registered validators
-	var wg sync.WaitGroup
 	clt := d.DefaultWallet().Client
 
 	// get the initial validators
 	expectedValidators := d.AccountsFromNodes(d.Nodes()...)
 
-	type validatorData struct {
-		wallet                    *mock.Wallet
-		implicitAccountOutputData *mock.OutputData
-	}
-
 	validatorCount := 50
-	validatorDataList := make([]validatorData, validatorCount)
-
-	for i := range validatorCount {
-		wg.Add(1)
-
-		// first create all implicit accounts in parallel
-		go func(validatorNr int) {
-			defer wg.Done()
 
-			// create implicit accounts for every validator
-			wallet, implicitAccountOutputData := d.CreateImplicitAccount(ctx)
-
-			validatorDataList[validatorNr] = validatorData{
-				wallet:                    wallet,
-				implicitAccountOutputData: implicitAccountOutputData,
-			}
-		}(i)
-	}
-
-	// wait until all implicit accounts are created
-	wg.Wait()
+	implicitAccountsValidators := d.CreateImplicitAccounts(ctx, validatorCount)
 
 	blockIssuance, err := clt.BlockIssuance(ctx)
 	require.NoError(t, err)
@@ -324,7 +299,7 @@ func Test_ValidatorsAPI(t *testing.T) {
 	stakingStartEpoch := d.DefaultWallet().StakingStartEpochFromSlot(latestCommitmentSlot)
 
 	// create a new wait group for the next step
-	wg = sync.WaitGroup{}
+	var wg sync.WaitGroup
 
 	// create accounts with staking feature and issue candidacy payload
 	for i := range validatorCount {
@@ -334,13 +309,12 @@ func Test_ValidatorsAPI(t *testing.T) {
 			defer wg.Done()
 
 			// create account with staking feature for every validator
-			accountData := d.CreateAccountFromImplicitAccount(validatorDataList[validatorNr].wallet,
-				validatorDataList[validatorNr].implicitAccountOutputData,
+			accountWithWallet := d.CreateAccountFromImplicitAccount(implicitAccountsValidators[validatorNr],
 				blockIssuance,
 				dockertestframework.WithStakingFeature(100, 1, stakingStartEpoch),
 			)
 
-			expectedValidators = append(expectedValidators, accountData.Address.Bech32(hrp))
+			expectedValidators = append(expectedValidators, accountWithWallet.Account().Address.Bech32(hrp))
 		}(i)
 	}
 	wg.Wait()

tools/docker-network/tests/committeerotation_test.go

@@ -217,21 +217,20 @@ func Test_Staking(t *testing.T) {
 	t.Cleanup(cancel)
 
 	// create implicit account for the validator
-	wallet, implicitAccountOutputData := d.CreateImplicitAccount(ctx)
+	implicitAccount := d.CreateImplicitAccount(ctx)
 
-	blockIssuance := wallet.GetNewBlockIssuanceResponse()
+	blockIssuance := implicitAccount.Wallet().GetNewBlockIssuanceResponse()
 
 	latestCommitmentSlot := blockIssuance.LatestCommitment.Slot
 	stakingStartEpoch := d.DefaultWallet().StakingStartEpochFromSlot(latestCommitmentSlot)
 
 	// create account with staking feature for the validator
-	accountData := d.CreateAccountFromImplicitAccount(wallet,
-		implicitAccountOutputData,
+	accountWithWallet := d.CreateAccountFromImplicitAccount(implicitAccount,
 		blockIssuance,
 		dockertestframework.WithStakingFeature(100, 1, stakingStartEpoch),
 	)
 
-	d.AssertValidatorExists(accountData.Address)
+	d.AssertValidatorExists(accountWithWallet.Account().Address)
 }
 
 // Test_Delegation tests if committee changed due to delegation.

tools/docker-network/tests/dockertestframework/accounts.go

@@ -5,6 +5,7 @@ package dockertestframework
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	"github.com/stretchr/testify/require"
 
@@ -59,43 +60,82 @@ func (d *DockerTestFramework) CheckAccountStatus(ctx context.Context, blkID iota
 	require.NoError(d.Testing, err)
 }
 
+type ImplicitAccount struct {
+	wallet     *mock.Wallet
+	outputData *mock.OutputData
+}
+
+func (i *ImplicitAccount) Wallet() *mock.Wallet {
+	return i.wallet
+}
+
+func (i *ImplicitAccount) OutputData() *mock.OutputData {
+	return i.outputData
+}
+
 // CreateImplicitAccount requests faucet funds and creates an implicit account. It already wait until the transaction is committed and the created account is useable.
-func (d *DockerTestFramework) CreateImplicitAccount(ctx context.Context) (*mock.Wallet, *mock.OutputData) {
-	newWallet := mock.NewWallet(d.Testing, "", d.defaultWallet.Client, &DockerWalletClock{client: d.defaultWallet.Client})
-	implicitAccountOutputData := d.RequestFaucetFunds(ctx, newWallet, iotago.AddressImplicitAccountCreation)
+func (d *DockerTestFramework) CreateImplicitAccount(ctx context.Context) *ImplicitAccount {
+	wallet := mock.NewWallet(d.Testing, "", d.defaultWallet.Client, &DockerWalletClock{client: d.defaultWallet.Client})
+	outputData := d.RequestFaucetFunds(ctx, wallet, iotago.AddressImplicitAccountCreation)
 
-	accountID := iotago.AccountIDFromOutputID(implicitAccountOutputData.ID)
+	accountID := iotago.AccountIDFromOutputID(outputData.ID)
 	accountAddress, ok := accountID.ToAddress().(*iotago.AccountAddress)
 	require.True(d.Testing, ok)
 
 	// make sure an implicit account is committed
-	d.CheckAccountStatus(ctx, iotago.EmptyBlockID, implicitAccountOutputData.ID.TransactionID(), implicitAccountOutputData.ID, accountAddress)
+	d.CheckAccountStatus(ctx, iotago.EmptyBlockID, outputData.ID.TransactionID(), outputData.ID, accountAddress)
 
 	// update the wallet with the new account data
-	newWallet.SetBlockIssuer(&mock.AccountData{
+	wallet.SetBlockIssuer(&mock.AccountData{
 		ID:           accountID,
 		Address:      accountAddress,
-		OutputID:     implicitAccountOutputData.ID,
-		AddressIndex: implicitAccountOutputData.AddressIndex,
+		OutputID:     outputData.ID,
+		AddressIndex: outputData.AddressIndex,
 	})
 
-	return newWallet, implicitAccountOutputData
+	return &ImplicitAccount{
+		wallet:     wallet,
+		outputData: outputData,
+	}
+}
+
+func (d *DockerTestFramework) CreateImplicitAccounts(ctx context.Context, count int) []*ImplicitAccount {
+	var wg sync.WaitGroup
+
+	implicitAccounts := make([]*ImplicitAccount, count)
+
+	for i := range count {
+		wg.Add(1)
+
+		// first create all implicit accounts in parallel
+		go func(nr int) {
+			defer wg.Done()
+
+			// create implicit accounts
+			implicitAccounts[nr] = d.CreateImplicitAccount(ctx)
+		}(i)
+	}
+
+	// wait until all implicit accounts are created
+	wg.Wait()
+
+	return implicitAccounts
 }
 
 // TransitionImplicitAccountToAccountOutputBlock consumes the given implicit account, then build the account transition block with the given account output options.
-func (d *DockerTestFramework) TransitionImplicitAccountToAccountOutputBlock(accountWallet *mock.Wallet, implicitAccountOutputData *mock.OutputData, blockIssuance *api.IssuanceBlockHeaderResponse, opts ...options.Option[builder.AccountOutputBuilder]) (*mock.AccountData, *iotago.SignedTransaction, *iotago.Block) {
+func (d *DockerTestFramework) TransitionImplicitAccountToAccountOutputBlock(implicitAccount *ImplicitAccount, blockIssuance *api.IssuanceBlockHeaderResponse, opts ...options.Option[builder.AccountOutputBuilder]) (*mock.AccountData, *iotago.SignedTransaction, *iotago.Block) {
 	ctx := context.TODO()
 
-	var implicitBlockIssuerKey iotago.BlockIssuerKey = iotago.Ed25519PublicKeyHashBlockIssuerKeyFromImplicitAccountCreationAddress(accountWallet.ImplicitAccountCreationAddress())
+	var implicitBlockIssuerKey iotago.BlockIssuerKey = iotago.Ed25519PublicKeyHashBlockIssuerKeyFromImplicitAccountCreationAddress(implicitAccount.Wallet().ImplicitAccountCreationAddress())
 	opts = append(opts, mock.WithBlockIssuerFeature(
 		iotago.NewBlockIssuerKeys(implicitBlockIssuerKey),
 		iotago.MaxSlotIndex,
 	))
 
-	signedTx := accountWallet.TransitionImplicitAccountToAccountOutputWithBlockIssuance("", []*mock.OutputData{implicitAccountOutputData}, blockIssuance, opts...)
+	signedTx := implicitAccount.Wallet().TransitionImplicitAccountToAccountOutputWithBlockIssuance("", []*mock.OutputData{implicitAccount.OutputData()}, blockIssuance, opts...)
 
 	// The account transition block should be issued by the implicit account block issuer key.
-	block, err := accountWallet.CreateBasicBlock(ctx, "", mock.WithPayload(signedTx))
+	block, err := implicitAccount.Wallet().CreateBasicBlock(ctx, "", mock.WithPayload(signedTx))
 	require.NoError(d.Testing, err)
 	accOutputID := iotago.OutputIDFromTransactionIDAndIndex(signedTx.Transaction.MustID(), 0)
 	accOutput := signedTx.Transaction.Outputs[0].(*iotago.AccountOutput)
@@ -106,47 +146,47 @@ func (d *DockerTestFramework) TransitionImplicitAccountToAccountOutputBlock(acco
 		Address:      accAddress,
 		Output:       accOutput,
 		OutputID:     accOutputID,
-		AddressIndex: implicitAccountOutputData.AddressIndex,
+		AddressIndex: implicitAccount.OutputData().AddressIndex,
 	}
 
 	return accountOutputData, signedTx, block.ProtocolBlock()
 }
 
 // CreateAccountFromImplicitAccount transitions an account from the given implicit one to full one, it already wait until the transaction is committed and the created account is useable.
-func (d *DockerTestFramework) CreateAccountFromImplicitAccount(accountWallet *mock.Wallet, implicitAccountOutputData *mock.OutputData, blockIssuance *api.IssuanceBlockHeaderResponse, opts ...options.Option[builder.AccountOutputBuilder]) *mock.AccountData {
+func (d *DockerTestFramework) CreateAccountFromImplicitAccount(implicitAccount *ImplicitAccount, blockIssuance *api.IssuanceBlockHeaderResponse, opts ...options.Option[builder.AccountOutputBuilder]) *mock.AccountWithWallet {
 	ctx := context.TODO()
 
-	accountData, signedTx, block := d.TransitionImplicitAccountToAccountOutputBlock(accountWallet, implicitAccountOutputData, blockIssuance, opts...)
+	accountData, signedTx, block := d.TransitionImplicitAccountToAccountOutputBlock(implicitAccount, blockIssuance, opts...)
 
 	d.SubmitBlock(ctx, block)
 	d.CheckAccountStatus(ctx, block.MustID(), signedTx.Transaction.MustID(), accountData.OutputID, accountData.Address, true)
 
 	// update the wallet with the new account data
-	accountWallet.SetBlockIssuer(accountData)
+	implicitAccount.Wallet().SetBlockIssuer(accountData)
 
-	fmt.Printf("Account created, Bech addr: %s\n", accountData.Address.Bech32(accountWallet.Client.CommittedAPI().ProtocolParameters().Bech32HRP()))
+	fmt.Printf("Account created, Bech addr: %s\n", accountData.Address.Bech32(implicitAccount.Wallet().Client.CommittedAPI().ProtocolParameters().Bech32HRP()))
 
-	return accountWallet.Account(accountData.ID)
+	return mock.NewAccountWithWallet(implicitAccount.Wallet().Account(accountData.ID), implicitAccount.Wallet())
 }
 
 // CreateAccountFromFaucet creates a new account by requesting faucet funds to an implicit account address and then transitioning the new output to a full account output.
 // It already waits until the transaction is committed and the created account is useable.
 func (d *DockerTestFramework) CreateAccountFromFaucet() (*mock.Wallet, *mock.AccountData) {
 	ctx := context.TODO()
 
-	newWallet, implicitAccountOutputData := d.CreateImplicitAccount(ctx)
+	implicitAccount := d.CreateImplicitAccount(ctx)
 
-	accountData, signedTx, block := d.TransitionImplicitAccountToAccountOutputBlock(newWallet, implicitAccountOutputData, d.defaultWallet.GetNewBlockIssuanceResponse())
+	accountData, signedTx, block := d.TransitionImplicitAccountToAccountOutputBlock(implicitAccount, d.defaultWallet.GetNewBlockIssuanceResponse())
 
 	d.SubmitBlock(ctx, block)
 	d.CheckAccountStatus(ctx, block.MustID(), signedTx.Transaction.MustID(), accountData.OutputID, accountData.Address, true)
 
 	// update the wallet with the new account data
-	newWallet.SetBlockIssuer(accountData)
+	implicitAccount.Wallet().SetBlockIssuer(accountData)
 
-	fmt.Printf("Account created, Bech addr: %s\n", accountData.Address.Bech32(newWallet.Client.CommittedAPI().ProtocolParameters().Bech32HRP()))
+	fmt.Printf("Account created, Bech addr: %s\n", accountData.Address.Bech32(implicitAccount.Wallet().Client.CommittedAPI().ProtocolParameters().Bech32HRP()))
 
-	return newWallet, newWallet.Account(accountData.ID)
+	return implicitAccount.Wallet(), implicitAccount.Wallet().Account(accountData.ID)
 }
 
 // CreateNativeToken request faucet funds then use it to create native token for the account, and returns the updated Account.

tools/docker-network/tests/eventapi_test.go

@@ -243,15 +243,15 @@ func test_AccountTransactionBlocks(t *testing.T, e *dockertestframework.EventAPI
 	// implicit account transition
 	{
 		// create an implicit account by requesting faucet funds
-		wallet, implicitAccountOutputData := e.DockerTestFramework().CreateImplicitAccount(ctx)
+		implicitAccount := e.DockerTestFramework().CreateImplicitAccount(ctx)
 
 		// prepare account transition block
-		accountData, _, blk := e.DockerTestFramework().TransitionImplicitAccountToAccountOutputBlock(wallet, implicitAccountOutputData, wallet.GetNewBlockIssuanceResponse())
+		accountData, _, blk := e.DockerTestFramework().TransitionImplicitAccountToAccountOutputBlock(implicitAccount, implicitAccount.Wallet().GetNewBlockIssuanceResponse())
 
 		expectedBlocks := map[string]*iotago.Block{
 			blk.MustID().ToHex(): blk,
 		}
-		accountOutputData := wallet.Output(accountData.OutputID)
+		accountOutputData := implicitAccount.Wallet().Output(accountData.OutputID)
 
 		assertions := []func(){
 			func() { e.AssertTransactionBlocks(ctx, eventClt, expectedBlocks) },