m/test/launch: make disk size configurable

This adds a flag to the launch-cluster command to specify the disk size 
of the VMs. This takes advantage of the previously added data partition 
growing feature.

Fixes: https://github.com/monogon-dev/monogon/issues/309
Change-Id: Iecf8d8c186af16dfa9ed0418ec96c51e58900052
Reviewed-on: https://review.monogon.dev/c/monogon/+/3351
Reviewed-by: Serge Bazanski <serge@monogon.tech>
Tested-by: Jenkins CI
diff --git a/metropolis/test/launch/cli/launch-cluster/main.go b/metropolis/test/launch/cli/launch-cluster/main.go
index cceebd1..3aeb319 100644
--- a/metropolis/test/launch/cli/launch-cluster/main.go
+++ b/metropolis/test/launch/cli/launch-cluster/main.go
@@ -72,7 +72,7 @@
 	})
 }
 
-func sizeFlagMiB(p *int, name string, usage string) {
+func memoryMiBFlag(p *int, name string, usage string) {
 	flag.Func(name, usage, func(val string) error {
 		multiplier := 1
 		switch {
@@ -91,6 +91,28 @@
 	})
 }
 
+func diskBytesFlag(p *uint64, name string, usage string) {
+	flag.Func(name, usage, func(val string) error {
+		var multiplier uint64
+		switch {
+		case strings.HasSuffix(val, "M"):
+			multiplier = 1024 * 1024
+		case strings.HasSuffix(val, "G"):
+			multiplier = 1024 * 1024 * 1024
+		case strings.HasSuffix(val, "T"):
+			multiplier = 1024 * 1024 * 1024 * 1024
+		default:
+			return errors.New("must have suffix M for MiB, G for GiB or T for TiB")
+		}
+		intVal, err := strconv.ParseUint(val[:len(val)-1], 10, 64)
+		if err != nil {
+			return err
+		}
+		*p = multiplier * intVal
+		return nil
+	})
+}
+
 func main() {
 	clusterConfig := cpb.ClusterConfiguration{}
 	opts := mlaunch.ClusterOptions{
@@ -104,7 +126,8 @@
 	flagdefs.StorageSecurityPolicyVar(flag.CommandLine, &clusterConfig.StorageSecurityPolicy, "storage-security", cpb.ClusterConfiguration_STORAGE_SECURITY_POLICY_NEEDS_INSECURE, "Storage security policy to set on cluster")
 	flag.IntVar(&opts.Node.CPUs, "cpu", 1, "Number of virtual CPUs of each node")
 	flag.IntVar(&opts.Node.ThreadsPerCPU, "threads-per-cpu", 1, "Number of threads per CPU")
-	sizeFlagMiB(&opts.Node.MemoryMiB, "ram", "RAM size of each node, with suffix M for MiB or G for GiB")
+	memoryMiBFlag(&opts.Node.MemoryMiB, "ram", "RAM size of each node, with suffix M for MiB or G for GiB")
+	diskBytesFlag(&opts.Node.DiskBytes, "disk", "Disk size of each node, with suffix M for MiB, G for GiB or T for TiB")
 	nodeSetFlag(&consensusMemberList, "consensus-member", "List of nodes which get the Consensus Member role. Example: 0,3-5")
 	nodeSetFlag(&kubernetesControllerList, "kubernetes-controller", "List of nodes which get the Kubernetes Controller role. Example: 0,3-5")
 	nodeSetFlag(&kubernetesWorkerList, "kubernetes-worker", "List of nodes which get the Kubernetes Worker role. Example: 0,3-5")