[antithesis] Refactor compose config generation to simplify reuse

Author: maru-ava

tests/antithesis/avalanchego/gencomposeconfig/main.go

@@ -4,9 +4,7 @@
 package main
 
 import (
-	"fmt"
 	"log"
-	"os"
 
 	"github.com/ava-labs/avalanchego/tests/antithesis"
 	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
@@ -16,22 +14,8 @@ const baseImageName = "antithesis-avalanchego"
 
 // Creates docker-compose.yml and its associated volumes in the target path.
 func main() {
-	targetPath := os.Getenv("TARGET_PATH")
-	if len(targetPath) == 0 {
-		log.Fatal("TARGET_PATH environment variable not set")
-	}
-
-	imageTag := os.Getenv("IMAGE_TAG")
-	if len(imageTag) == 0 {
-		log.Fatal("IMAGE_TAG environment variable not set")
-	}
-
-	nodeImageName := fmt.Sprintf("%s-node:%s", baseImageName, imageTag)
-	workloadImageName := fmt.Sprintf("%s-workload:%s", baseImageName, imageTag)
-
 	network := tmpnet.LocalNetworkOrPanic()
-	err := antithesis.GenerateComposeConfig(network, nodeImageName, workloadImageName, targetPath)
-	if err != nil {
-		log.Fatalf("failed to generate config for docker-compose: %s", err)
+	if err := antithesis.GenerateComposeConfig(network, baseImageName); err != nil {
+		log.Fatalf("failed to generate compose config: %v", err)
 	}
 }

tests/antithesis/compose.go

@@ -4,6 +4,7 @@
 package antithesis
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -22,9 +23,56 @@ import (
 
 const bootstrapIndex = 0
 
+// Creates docker-compose configuration for an antithesis test
+// setup. Configuration is via env vars to simplify usage by main entrypoints. If
+// the provided network includes a subnet, the initial DB state for the subnet
+// will be created and written to the target path.
+func GenerateComposeConfig(network *tmpnet.Network, baseImageName string) error {
+	targetPath := os.Getenv("TARGET_PATH")
+	if len(targetPath) == 0 {
+		return errors.New("TARGET_PATH environment variable not set")
+	}
+
+	imageTag := os.Getenv("IMAGE_TAG")
+	if len(imageTag) == 0 {
+		return errors.New("IMAGE_TAG environment variable not set")
+	}
+
+	// Subnet testing requires creating an initial db state for the bootstrap node
+	if len(network.Subnets) > 0 {
+		avalancheGoPath := os.Getenv("AVALANCHEGO_PATH")
+		if len(avalancheGoPath) == 0 {
+			return errors.New("AVALANCHEGO_PATH environment variable not set")
+		}
+
+		pluginDir := os.Getenv("AVALANCHEGO_PLUGIN_DIR")
+		if len(pluginDir) == 0 {
+			return errors.New("AVALANCHEGO_PLUGIN_DIR environment variable not set")
+		}
+
+		bootstrapVolumePath, err := getBootstrapVolumePath(targetPath)
+		if err != nil {
+			return fmt.Errorf("failed to get bootstrap volume path: %w", err)
+		}
+
+		if err := initBootstrapDB(network, avalancheGoPath, pluginDir, bootstrapVolumePath); err != nil {
+			return fmt.Errorf("failed to initialize db volumes: %w", err)
+		}
+	}
+
+	nodeImageName := fmt.Sprintf("%s-node:%s", baseImageName, imageTag)
+	workloadImageName := fmt.Sprintf("%s-workload:%s", baseImageName, imageTag)
+
+	if err := initComposeConfig(network, nodeImageName, workloadImageName, targetPath); err != nil {
+		return fmt.Errorf("failed to generate compose config: %w", err)
+	}
+
+	return nil
+}
+
 // Initialize the given path with the docker-compose configuration (compose file and
 // volumes) needed for an Antithesis test setup.
-func GenerateComposeConfig(
+func initComposeConfig(
 	network *tmpnet.Network,
 	nodeImageName string,
 	workloadImageName string,

tests/antithesis/init_db.go

@@ -16,7 +16,7 @@ import (
 )
 
 // Given a path, compose the expected path of the bootstrap node's docker compose db volume.
-func GetBootstrapVolumePath(targetPath string) (string, error) {
+func getBootstrapVolumePath(targetPath string) (string, error) {
 	absPath, err := filepath.Abs(targetPath)
 	if err != nil {
 		return "", fmt.Errorf("failed to convert target path to absolute path: %w", err)
@@ -27,7 +27,7 @@ func GetBootstrapVolumePath(targetPath string) (string, error) {
 // Bootstraps a local process-based network, creates its subnets and chains, and copies
 // the resulting db state from one of the nodes to the provided path. The path will be
 // created if it does not already exist.
-func InitBootstrapDB(network *tmpnet.Network, avalancheGoPath string, pluginDir string, destPath string) error {
+func initBootstrapDB(network *tmpnet.Network, avalancheGoPath string, pluginDir string, destPath string) error {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*2)
 	defer cancel()
 	if err := tmpnet.BootstrapNewNetwork(

tests/antithesis/xsvm/gencomposeconfig/main.go

@@ -4,9 +4,7 @@
 package main
 
 import (
-	"fmt"
 	"log"
-	"os"
 
 	"github.com/ava-labs/avalanchego/genesis"
 	"github.com/ava-labs/avalanchego/tests/antithesis"
@@ -18,45 +16,11 @@ const baseImageName = "antithesis-xsvm"
 
 // Creates docker-compose.yml and its associated volumes in the target path.
 func main() {
-	avalancheGoPath := os.Getenv("AVALANCHEGO_PATH")
-	if len(avalancheGoPath) == 0 {
-		log.Fatal("AVALANCHEGO_PATH environment variable not set")
-	}
-
-	pluginDir := os.Getenv("AVALANCHEGO_PLUGIN_DIR")
-	if len(pluginDir) == 0 {
-		log.Fatal("AVALANCHEGO_PLUGIN_DIR environment variable not set")
-	}
-
-	targetPath := os.Getenv("TARGET_PATH")
-	if len(targetPath) == 0 {
-		log.Fatal("TARGET_PATH environment variable not set")
-	}
-
-	imageTag := os.Getenv("IMAGE_TAG")
-	if len(imageTag) == 0 {
-		log.Fatal("IMAGE_TAG environment variable not set")
-	}
-
-	nodeImageName := fmt.Sprintf("%s-node:%s", baseImageName, imageTag)
-	workloadImageName := fmt.Sprintf("%s-workload:%s", baseImageName, imageTag)
-
-	// Create a network with an xsvm subnet
 	network := tmpnet.LocalNetworkOrPanic()
 	network.Subnets = []*tmpnet.Subnet{
 		subnet.NewXSVMOrPanic("xsvm", genesis.VMRQKey, network.Nodes...),
 	}
-
-	bootstrapVolumePath, err := antithesis.GetBootstrapVolumePath(targetPath)
-	if err != nil {
-		log.Fatalf("failed to get bootstrap volume path: %v", err)
-	}
-
-	if err := antithesis.InitBootstrapDB(network, avalancheGoPath, pluginDir, bootstrapVolumePath); err != nil {
-		log.Fatalf("failed to initialize db volumes: %v", err)
-	}
-
-	if err := antithesis.GenerateComposeConfig(network, nodeImageName, workloadImageName, targetPath); err != nil {
-		log.Fatalf("failed to generate config for docker-compose: %v", err)
+	if err := antithesis.GenerateComposeConfig(network, baseImageName); err != nil {
+		log.Fatalf("failed to generate compose config: %v", err)
 	}
 }