github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/cmd/cortex/main_test.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"flag"
     6  	"io"
     7  	"io/ioutil"
     8  	"os"
     9  	"strings"
    10  	"sync"
    11  	"testing"
    12  
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  func TestFlagParsing(t *testing.T) {
    18  	for name, tc := range map[string]struct {
    19  		arguments      []string
    20  		yaml           string
    21  		stdoutMessage  string // string that must be included in stdout
    22  		stderrMessage  string // string that must be included in stderr
    23  		stdoutExcluded string // string that must NOT be included in stdout
    24  		stderrExcluded string // string that must NOT be included in stderr
    25  	}{
    26  		"help": {
    27  			arguments:      []string{"-h"},
    28  			stdoutMessage:  "Usage of", // Usage must be on stdout, not stderr.
    29  			stderrExcluded: "Usage of",
    30  		},
    31  
    32  		"unknown flag": {
    33  			arguments:      []string{"-unknown.flag"},
    34  			stderrMessage:  "Run with -help to get list of available parameters",
    35  			stdoutExcluded: "Usage of", // No usage description on unknown flag.
    36  			stderrExcluded: "Usage of",
    37  		},
    38  
    39  		"new flag, with config": {
    40  			arguments:     []string{"-mem-ballast-size-bytes=100000"},
    41  			yaml:          "target: ingester",
    42  			stdoutMessage: "target: ingester",
    43  		},
    44  
    45  		"default values": {
    46  			stdoutMessage: "target: all\n",
    47  		},
    48  
    49  		"config": {
    50  			yaml:          "target: ingester",
    51  			stdoutMessage: "target: ingester\n",
    52  		},
    53  
    54  		"config with expand-env": {
    55  			arguments:     []string{"-config.expand-env"},
    56  			yaml:          "target: $TARGET",
    57  			stdoutMessage: "target: ingester\n",
    58  		},
    59  
    60  		"config with arguments override": {
    61  			yaml:          "target: ingester",
    62  			arguments:     []string{"-target=distributor"},
    63  			stdoutMessage: "target: distributor\n",
    64  		},
    65  
    66  		"user visible module listing": {
    67  			arguments:      []string{"-modules"},
    68  			stdoutMessage:  "ingester *\n",
    69  			stderrExcluded: "ingester\n",
    70  		},
    71  
    72  		"user visible module listing flag take precedence over target flag": {
    73  			arguments:      []string{"-modules", "-target=blah"},
    74  			stdoutMessage:  "ingester *\n",
    75  			stderrExcluded: "ingester\n",
    76  		},
    77  
    78  		"root level configuration option specified as an empty node in YAML": {
    79  			yaml:          "querier:",
    80  			stderrMessage: "the Querier configuration in YAML has been specified as an empty YAML node",
    81  		},
    82  
    83  		"version": {
    84  			arguments:     []string{"-version"},
    85  			stdoutMessage: "Cortex, version",
    86  		},
    87  
    88  		// we cannot test the happy path, as cortex would then fully start
    89  	} {
    90  		t.Run(name, func(t *testing.T) {
    91  			_ = os.Setenv("TARGET", "ingester")
    92  			testSingle(t, tc.arguments, tc.yaml, []byte(tc.stdoutMessage), []byte(tc.stderrMessage), []byte(tc.stdoutExcluded), []byte(tc.stderrExcluded))
    93  		})
    94  	}
    95  }
    96  
    97  func testSingle(t *testing.T, arguments []string, yaml string, stdoutMessage, stderrMessage, stdoutExcluded, stderrExcluded []byte) {
    98  	t.Helper()
    99  	oldArgs, oldStdout, oldStderr, oldTestMode := os.Args, os.Stdout, os.Stderr, testMode
   100  	restored := false
   101  	restoreIfNeeded := func() {
   102  		if restored {
   103  			return
   104  		}
   105  		os.Stdout = oldStdout
   106  		os.Stderr = oldStderr
   107  		os.Args = oldArgs
   108  		testMode = oldTestMode
   109  		restored = true
   110  	}
   111  	defer restoreIfNeeded()
   112  
   113  	if yaml != "" {
   114  		tempFile, err := ioutil.TempFile("", "test")
   115  		require.NoError(t, err)
   116  
   117  		defer func() {
   118  			require.NoError(t, tempFile.Close())
   119  			require.NoError(t, os.Remove(tempFile.Name()))
   120  		}()
   121  
   122  		_, err = tempFile.WriteString(yaml)
   123  		require.NoError(t, err)
   124  
   125  		arguments = append(arguments, "-"+configFileOption, tempFile.Name())
   126  	}
   127  
   128  	arguments = append([]string{"./cortex"}, arguments...)
   129  
   130  	testMode = true
   131  	os.Args = arguments
   132  	co := captureOutput(t)
   133  
   134  	// reset default flags
   135  	flag.CommandLine = flag.NewFlagSet(arguments[0], flag.ExitOnError)
   136  
   137  	main()
   138  
   139  	stdout, stderr := co.Done()
   140  
   141  	// Restore stdout and stderr before reporting errors to make them visible.
   142  	restoreIfNeeded()
   143  	if !bytes.Contains(stdout, stdoutMessage) {
   144  		t.Errorf("Expected on stdout: %q, stdout: %s\n", stdoutMessage, stdout)
   145  	}
   146  	if !bytes.Contains(stderr, stderrMessage) {
   147  		t.Errorf("Expected on stderr: %q, stderr: %s\n", stderrMessage, stderr)
   148  	}
   149  	if len(stdoutExcluded) > 0 && bytes.Contains(stdout, stdoutExcluded) {
   150  		t.Errorf("Unexpected output on stdout: %q, stdout: %s\n", stdoutExcluded, stdout)
   151  	}
   152  	if len(stderrExcluded) > 0 && bytes.Contains(stderr, stderrExcluded) {
   153  		t.Errorf("Unexpected output on stderr: %q, stderr: %s\n", stderrExcluded, stderr)
   154  	}
   155  }
   156  
   157  type capturedOutput struct {
   158  	stdoutBuf bytes.Buffer
   159  	stderrBuf bytes.Buffer
   160  
   161  	wg                         sync.WaitGroup
   162  	stdoutReader, stdoutWriter *os.File
   163  	stderrReader, stderrWriter *os.File
   164  }
   165  
   166  func captureOutput(t *testing.T) *capturedOutput {
   167  	stdoutR, stdoutW, err := os.Pipe()
   168  	require.NoError(t, err)
   169  	os.Stdout = stdoutW
   170  
   171  	stderrR, stderrW, err := os.Pipe()
   172  	require.NoError(t, err)
   173  	os.Stderr = stderrW
   174  
   175  	co := &capturedOutput{
   176  		stdoutReader: stdoutR,
   177  		stdoutWriter: stdoutW,
   178  		stderrReader: stderrR,
   179  		stderrWriter: stderrW,
   180  	}
   181  	co.wg.Add(1)
   182  	go func() {
   183  		defer co.wg.Done()
   184  		io.Copy(&co.stdoutBuf, stdoutR)
   185  	}()
   186  
   187  	co.wg.Add(1)
   188  	go func() {
   189  		defer co.wg.Done()
   190  		io.Copy(&co.stderrBuf, stderrR)
   191  	}()
   192  
   193  	return co
   194  }
   195  
   196  func (co *capturedOutput) Done() (stdout []byte, stderr []byte) {
   197  	// we need to close writers for readers to stop
   198  	_ = co.stdoutWriter.Close()
   199  	_ = co.stderrWriter.Close()
   200  
   201  	co.wg.Wait()
   202  
   203  	return co.stdoutBuf.Bytes(), co.stderrBuf.Bytes()
   204  }
   205  
   206  func TestExpandEnv(t *testing.T) {
   207  	var tests = []struct {
   208  		in  string
   209  		out string
   210  	}{
   211  		// Environment variables can be specified as ${env} or $env.
   212  		{"x$y", "xy"},
   213  		{"x${y}", "xy"},
   214  
   215  		// Environment variables are case-sensitive. Neither are replaced.
   216  		{"x$Y", "x"},
   217  		{"x${Y}", "x"},
   218  
   219  		// Defaults can only be specified when using braces.
   220  		{"x${Z:D}", "xD"},
   221  		{"x${Z:A B C D}", "xA B C D"}, // Spaces are allowed in the default.
   222  		{"x${Z:}", "x"},
   223  
   224  		// Defaults don't work unless braces are used.
   225  		{"x$y:D", "xy:D"},
   226  	}
   227  
   228  	for _, test := range tests {
   229  		test := test
   230  		t.Run(test.in, func(t *testing.T) {
   231  			_ = os.Setenv("y", "y")
   232  			output := expandEnv([]byte(test.in))
   233  			assert.Equal(t, test.out, string(output), "Input: %s", test.in)
   234  		})
   235  	}
   236  }
   237  
   238  func TestParseConfigFileParameter(t *testing.T) {
   239  	var tests = []struct {
   240  		args       string
   241  		configFile string
   242  		expandENV  bool
   243  	}{
   244  		{"", "", false},
   245  		{"--foo", "", false},
   246  		{"-f -a", "", false},
   247  
   248  		{"--config.file=foo", "foo", false},
   249  		{"--config.file foo", "foo", false},
   250  		{"--config.file=foo --config.expand-env", "foo", true},
   251  		{"--config.expand-env --config.file=foo", "foo", true},
   252  
   253  		{"--opt1 --config.file=foo", "foo", false},
   254  		{"--opt1 --config.file foo", "foo", false},
   255  		{"--opt1 --config.file=foo --config.expand-env", "foo", true},
   256  		{"--opt1 --config.expand-env --config.file=foo", "foo", true},
   257  
   258  		{"--config.file=foo --opt1", "foo", false},
   259  		{"--config.file foo --opt1", "foo", false},
   260  		{"--config.file=foo --config.expand-env --opt1", "foo", true},
   261  		{"--config.expand-env --config.file=foo --opt1", "foo", true},
   262  
   263  		{"--config.file=foo --opt1 --config.expand-env", "foo", true},
   264  		{"--config.expand-env --opt1 --config.file=foo", "foo", true},
   265  	}
   266  	for _, test := range tests {
   267  		test := test
   268  		t.Run(test.args, func(t *testing.T) {
   269  			args := strings.Split(test.args, " ")
   270  			configFile, expandENV := parseConfigFileParameter(args)
   271  			assert.Equal(t, test.configFile, configFile)
   272  			assert.Equal(t, test.expandENV, expandENV)
   273  		})
   274  	}
   275  }