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 }