Refactor docker test framework.

Author: piotrm50

tools/docker-network/tests/committeerotation_test.go

@@ -146,13 +146,13 @@ func Test_NoCandidacyPayload(t *testing.T) {
 		))
 	defer d.Stop()
 
-	d.AddValidatorNode("V1", "docker-network-inx-validator-1-1", "http://localhost:8050", "rms1pzg8cqhfxqhq7pt37y8cs4v5u4kcc48lquy2k73ehsdhf5ukhya3y5rx2w6")
-	d.AddValidatorNode("V2", "docker-network-inx-validator-2-1", "http://localhost:8060", "rms1pqm4xk8e9ny5w5rxjkvtp249tfhlwvcshyr3pc0665jvp7g3hc875k538hl")
-	d.AddValidatorNode("V3", "docker-network-inx-validator-3-1", "http://localhost:8070", "rms1pp4wuuz0y42caz48vv876qfpmffswsvg40zz8v79sy8cp0jfxm4kunflcgt")
-	d.AddValidatorNode("V4", "docker-network-inx-validator-4-1", "http://localhost:8040", "rms1pr8cxs3dzu9xh4cduff4dd4cxdthpjkpwmz2244f75m0urslrsvtsshrrjw")
+	d.AddValidatorNode("V1", "docker-network-inx-validator-1-1", "http://localhost:8050", "rms1pzg8cqhfxqhq7pt37y8cs4v5u4kcc48lquy2k73ehsdhf5ukhya3y5rx2w6", false)
+	d.AddValidatorNode("V2", "docker-network-inx-validator-2-1", "http://localhost:8060", "rms1pqm4xk8e9ny5w5rxjkvtp249tfhlwvcshyr3pc0665jvp7g3hc875k538hl", false)
+	d.AddValidatorNode("V3", "docker-network-inx-validator-3-1", "http://localhost:8070", "rms1pp4wuuz0y42caz48vv876qfpmffswsvg40zz8v79sy8cp0jfxm4kunflcgt", false)
+	d.AddValidatorNode("V4", "docker-network-inx-validator-4-1", "http://localhost:8040", "rms1pr8cxs3dzu9xh4cduff4dd4cxdthpjkpwmz2244f75m0urslrsvtsshrrjw", false)
 	d.AddNode("node5", "docker-network-node-5-1", "http://localhost:8090")
 
-	err := d.Run(map[string]bool{"V1": false, "V2": false, "V3": false, "V4": false})
+	err := d.Run()
 	require.NoError(t, err)
 
 	err = d.WaitUntilSync()
@@ -179,7 +179,7 @@ func Test_NoCandidacyPayload(t *testing.T) {
 	})
 
 	// Start issuing candidacy payloads for 3 validators, and check if committee size is 3
-	d.SetIssueCandidacyPayload(map[string]bool{"V1": true, "V2": true, "V3": true, "V4": false})
+	d.StartIssueCandidacyPayload("V1", "V2", "V3")
 	d.AssertCommittee(currentEpoch+4, d.AccountsFromNodes(d.Nodes("V1", "V2", "V3")...))
 }
 

tools/docker-network/tests/dockerframework.go

@@ -9,6 +9,7 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"slices"
 	"sort"
 	"strings"
 	"sync/atomic"
@@ -49,13 +50,14 @@ var (
 )
 
 type Node struct {
-	Name                 string
-	ContainerName        string
-	ClientURL            string
-	Client               *nodeclient.Client
-	AccountAddressBech32 string
-	ContainerConfigs     string
-	PrivateKey           string
+	Name                  string
+	ContainerName         string
+	ClientURL             string
+	Client                *nodeclient.Client
+	AccountAddressBech32  string
+	ContainerConfigs      string
+	PrivateKey            string
+	IssueCandidacyPayload bool
 }
 
 type Account struct {
@@ -116,31 +118,36 @@ func NewDockerTestFramework(t *testing.T, opts ...options.Option[DockerTestFrame
 	})
 }
 
-func (d *DockerTestFramework) Run(optIssueCandidacyPayloadMap ...map[string]bool) error {
+func (d *DockerTestFramework) DockerComposeUp(detach ...bool) error {
+	cmd := exec.Command("docker", "compose", "up")
+
+	if len(detach) > 0 && detach[0] {
+		cmd = exec.Command("docker", "compose", "up", "-d")
+	}
+
+	cmd.Env = os.Environ()
+	for _, node := range d.nodes {
+		cmd.Env = append(cmd.Env, fmt.Sprintf("ISSUE_CANDIDACY_PAYLOAD_%s=%t", node.Name, node.IssueCandidacyPayload))
+	}
+
+	var out strings.Builder
+	cmd.Stderr = &out
+	err := cmd.Run()
+	if err != nil {
+		fmt.Println("Docker compose up failed with error:", err, ":", out.String())
+	}
+
+	return err
+}
+
+func (d *DockerTestFramework) Run() error {
 	ch := make(chan error)
 	stopCh := make(chan struct{})
 	defer close(ch)
 	defer close(stopCh)
 
-	issueCandidacyPayloadMap := make(map[string]bool)
-	if len(optIssueCandidacyPayloadMap) > 0 {
-		issueCandidacyPayloadMap = optIssueCandidacyPayloadMap[0]
-	}
-
 	go func() {
-		cmd := exec.Command("docker", "compose", "up")
-		cmd.Env = os.Environ()
-		for nodeName, issueCandidacyPayload := range issueCandidacyPayloadMap {
-			cmd.Env = append(cmd.Env, fmt.Sprintf("ISSUE_CANDIDACY_PAYLOAD_%s=%t", nodeName, issueCandidacyPayload))
-		}
-
-		var out strings.Builder
-		cmd.Stderr = &out
-		err := cmd.Run()
-
-		if err != nil {
-			fmt.Println("Docker compose up failed with error:", err, ":", out.String())
-		}
+		err := d.DockerComposeUp()
 
 		// make sure that the channel is not already closed
 		select {
@@ -222,12 +229,18 @@ func (d *DockerTestFramework) WaitUntilSync() error {
 	return nil
 }
 
-func (d *DockerTestFramework) AddValidatorNode(name string, containerName string, clientURL string, accAddrBech32 string) {
+func (d *DockerTestFramework) AddValidatorNode(name string, containerName string, clientURL string, accAddrBech32 string, optIssueCandidacyPayload ...bool) {
+	issueCandidacyPayload := true
+	if len(optIssueCandidacyPayload) > 0 {
+		issueCandidacyPayload = optIssueCandidacyPayload[0]
+	}
+
 	d.nodes[name] = &Node{
-		Name:                 name,
-		ContainerName:        containerName,
-		ClientURL:            clientURL,
-		AccountAddressBech32: accAddrBech32,
+		Name:                  name,
+		ContainerName:         containerName,
+		ClientURL:             clientURL,
+		AccountAddressBech32:  accAddrBech32,
+		IssueCandidacyPayload: issueCandidacyPayload,
 	}
 }
 
@@ -284,23 +297,36 @@ func (d *DockerTestFramework) AccountsFromNodes(nodes ...*Node) []string {
 	return accounts
 }
 
-func (d *DockerTestFramework) SetIssueCandidacyPayload(issueCandidacyPayloadMap map[string]bool) {
-	cmd := exec.Command("docker", "compose", "up", "-d")
+func (d *DockerTestFramework) StartIssueCandidacyPayload(nodeNames ...string) {
+	if len(nodeNames) == 0 {
+		return
+	}
 
-	cmd.Env = os.Environ()
+	for _, node := range d.nodes {
+		if slices.Contains(nodeNames, node.Name) {
+			node.IssueCandidacyPayload = true
+		}
+	}
+
+	d.DockerComposeUp(true)
+}
 
-	for nodeName, issueCandidacyPayload := range issueCandidacyPayloadMap {
-		cmd.Env = append(cmd.Env, fmt.Sprintf("ISSUE_CANDIDACY_PAYLOAD_%s=%t", nodeName, issueCandidacyPayload))
+func (d *DockerTestFramework) StopIssueCandidacyPayload(nodeNames ...string) {
+	if len(nodeNames) == 0 {
+		return
 	}
 
-	err := cmd.Run()
-	if err != nil {
-		fmt.Println("Docker compose up failed with error:", err)
+	for _, node := range d.nodes {
+		if slices.Contains(nodeNames, node.Name) {
+			node.IssueCandidacyPayload = false
+		}
 	}
+
+	d.DockerComposeUp(true)
 }
 
 func (d *DockerTestFramework) CreateAccount(opts ...options.Option[builder.AccountOutputBuilder]) *Account {
-	// create implicit account by requesting faucet funds
+	// create an implicit account by requesting faucet funds
 	ctx := context.TODO()
 	receiverAddr, implicitPrivateKey := d.getAddress(iotago.AddressImplicitAccountCreation)
 	implicitOutputID, implicitAccountOutput := d.RequestFaucetFunds(ctx, receiverAddr)
@@ -309,10 +335,10 @@ func (d *DockerTestFramework) CreateAccount(opts ...options.Option[builder.Accou
 	accountAddress, ok := accountID.ToAddress().(*iotago.AccountAddress)
 	require.True(d.Testing, ok)
 
-	// make sure implicit account is committed
+	// make sure an implicit account is committed
 	d.CheckAccountStatus(ctx, iotago.EmptyBlockID, implicitOutputID.TransactionID(), implicitOutputID, accountAddress)
 
-	// transition to full account with new Ed25519 address and staking feature
+	// transition to a full account with new Ed25519 address and staking feature
 	accEd25519Addr, accPrivateKey := d.getAddress(iotago.AddressEd25519)
 	accBlockIssuerKey := iotago.Ed25519PublicKeyHashBlockIssuerKeyFromPublicKey(accPrivateKey.Public().(ed25519.PublicKey))
 	accountOutput := options.Apply(builder.NewAccountOutputBuilder(accEd25519Addr, implicitAccountOutput.BaseTokenAmount()),