From dec668f75ae02e6f310f1893aa7b03786a37deda Mon Sep 17 00:00:00 2001
From: tsachiherman <24438559+tsachiherman@users.noreply.github.com>
Date: Fri, 6 Sep 2024 13:12:15 -0400
Subject: [PATCH] update

---
 .../avalanchego/gencomposeconfig/main.go          |  8 +++++---
 tests/antithesis/network.go                       | 15 +++++++++++++++
 tests/antithesis/xsvm/gencomposeconfig/main.go    |  7 +++++--
 3 files changed, 25 insertions(+), 5 deletions(-)
 create mode 100644 tests/antithesis/network.go

diff --git a/tests/antithesis/avalanchego/gencomposeconfig/main.go b/tests/antithesis/avalanchego/gencomposeconfig/main.go
index 7425a3df9bba..689375859c17 100644
--- a/tests/antithesis/avalanchego/gencomposeconfig/main.go
+++ b/tests/antithesis/avalanchego/gencomposeconfig/main.go
@@ -7,14 +7,16 @@ import (
 	"log"
 
 	"github.com/ava-labs/avalanchego/tests/antithesis"
-	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
 )
 
-const baseImageName = "antithesis-avalanchego"
+const (
+	baseImageName = "antithesis-avalanchego"
+	NumNodes      = 6
+)
 
 // Creates docker-compose.yml and its associated volumes in the target path.
 func main() {
-	network := tmpnet.LocalNetworkOrPanic()
+	network := antithesis.CreateNetwork(NumNodes)
 	if err := antithesis.GenerateComposeConfig(network, baseImageName, "" /* runtimePluginDir */); err != nil {
 		log.Fatalf("failed to generate compose config: %v", err)
 	}
diff --git a/tests/antithesis/network.go b/tests/antithesis/network.go
new file mode 100644
index 000000000000..7bdb335b7b2e
--- /dev/null
+++ b/tests/antithesis/network.go
@@ -0,0 +1,15 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package antithesis
+
+import (
+	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
+)
+
+// Creates a network with the given number of nodes.
+func CreateNetwork(nodesCount int) *tmpnet.Network {
+	network := tmpnet.NewDefaultNetwork("antithesis-network")
+	network.Nodes = tmpnet.NewNodesOrPanic(nodesCount)
+	return network
+}
diff --git a/tests/antithesis/xsvm/gencomposeconfig/main.go b/tests/antithesis/xsvm/gencomposeconfig/main.go
index e038daac0893..c3c89630f3c3 100644
--- a/tests/antithesis/xsvm/gencomposeconfig/main.go
+++ b/tests/antithesis/xsvm/gencomposeconfig/main.go
@@ -12,11 +12,14 @@ import (
 	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
 )
 
-const baseImageName = "antithesis-xsvm"
+const (
+	baseImageName = "antithesis-xsvm"
+	NumNodes      = 6
+)
 
 // Creates docker-compose.yml and its associated volumes in the target path.
 func main() {
-	network := tmpnet.LocalNetworkOrPanic()
+	network := antithesis.CreateNetwork(NumNodes)
 	network.Subnets = []*tmpnet.Subnet{
 		subnet.NewXSVMOrPanic("xsvm", genesis.VMRQKey, network.Nodes...),
 	}