github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/cmd/roachprod/vm/vm_test.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package vm
    12  
    13  import (
    14  	"strconv"
    15  	"testing"
    16  
    17  	"github.com/stretchr/testify/assert"
    18  )
    19  
    20  func TestZonePlacement(t *testing.T) {
    21  	for i, c := range []struct {
    22  		numZones, numNodes int
    23  		expected           []int
    24  	}{
    25  		{1, 1, []int{0}},
    26  		{1, 2, []int{0, 0}},
    27  		{2, 4, []int{0, 0, 1, 1}},
    28  		{2, 5, []int{0, 0, 1, 1, 0}},
    29  		{3, 11, []int{0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 1}},
    30  	} {
    31  		t.Run(strconv.Itoa(i), func(t *testing.T) {
    32  			assert.EqualValues(t, c.expected, ZonePlacement(c.numZones, c.numNodes))
    33  		})
    34  	}
    35  }
    36  
    37  func TestExpandZonesFlag(t *testing.T) {
    38  	for i, c := range []struct {
    39  		input, output []string
    40  		expErr        string
    41  	}{
    42  		{
    43  			input:  []string{"us-east1-b:3", "us-west2-c:2"},
    44  			output: []string{"us-east1-b", "us-east1-b", "us-east1-b", "us-west2-c", "us-west2-c"},
    45  		},
    46  		{
    47  			input:  []string{"us-east1-b:3", "us-west2-c"},
    48  			output: []string{"us-east1-b", "us-east1-b", "us-east1-b", "us-west2-c"},
    49  		},
    50  		{
    51  			input:  []string{"us-east1-b", "us-west2-c"},
    52  			output: []string{"us-east1-b", "us-west2-c"},
    53  		},
    54  		{
    55  			input:  []string{"us-east1-b", "us-west2-c:a2"},
    56  			expErr: "failed to parse",
    57  		},
    58  	} {
    59  		t.Run(strconv.Itoa(i), func(t *testing.T) {
    60  			expanded, err := ExpandZonesFlag(c.input)
    61  			if c.expErr != "" {
    62  				if assert.Error(t, err) {
    63  					assert.Regexp(t, c.expErr, err.Error())
    64  				}
    65  			} else {
    66  				assert.EqualValues(t, c.output, expanded)
    67  			}
    68  		})
    69  	}
    70  }