| package logrus |
| |
| import ( |
| "io/ioutil" |
| "log" |
| "os" |
| "os/exec" |
| "path/filepath" |
| "runtime" |
| "strings" |
| "testing" |
| "time" |
| ) |
| |
| func TestRegister(t *testing.T) { |
| current := len(handlers) |
| |
| var results []string |
| |
| h1 := func() { results = append(results, "first") } |
| h2 := func() { results = append(results, "second") } |
| |
| RegisterExitHandler(h1) |
| RegisterExitHandler(h2) |
| |
| if len(handlers) != current+2 { |
| t.Fatalf("expected %d handlers, got %d", current+2, len(handlers)) |
| } |
| |
| runHandlers() |
| |
| if len(results) != 2 { |
| t.Fatalf("expected 2 handlers to be run, ran %d", len(results)) |
| } |
| |
| if results[0] != "first" { |
| t.Fatal("expected handler h1 to be run first, but it wasn't") |
| } |
| |
| if results[1] != "second" { |
| t.Fatal("expected handler h2 to be run second, but it wasn't") |
| } |
| } |
| |
| func TestDefer(t *testing.T) { |
| current := len(handlers) |
| |
| var results []string |
| |
| h1 := func() { results = append(results, "first") } |
| h2 := func() { results = append(results, "second") } |
| |
| DeferExitHandler(h1) |
| DeferExitHandler(h2) |
| |
| if len(handlers) != current+2 { |
| t.Fatalf("expected %d handlers, got %d", current+2, len(handlers)) |
| } |
| |
| runHandlers() |
| |
| if len(results) != 2 { |
| t.Fatalf("expected 2 handlers to be run, ran %d", len(results)) |
| } |
| |
| if results[0] != "second" { |
| t.Fatal("expected handler h2 to be run first, but it wasn't") |
| } |
| |
| if results[1] != "first" { |
| t.Fatal("expected handler h1 to be run second, but it wasn't") |
| } |
| } |
| |
| func TestHandler(t *testing.T) { |
| testprog := testprogleader |
| testprog = append(testprog, getPackage()...) |
| testprog = append(testprog, testprogtrailer...) |
| tempDir, err := ioutil.TempDir("", "test_handler") |
| if err != nil { |
| log.Fatalf("can't create temp dir. %q", err) |
| } |
| defer os.RemoveAll(tempDir) |
| |
| gofile := filepath.Join(tempDir, "gofile.go") |
| if err := ioutil.WriteFile(gofile, testprog, 0666); err != nil { |
| t.Fatalf("can't create go file. %q", err) |
| } |
| |
| outfile := filepath.Join(tempDir, "outfile.out") |
| arg := time.Now().UTC().String() |
| err = exec.Command("go", "run", gofile, outfile, arg).Run() |
| if err == nil { |
| t.Fatalf("completed normally, should have failed") |
| } |
| |
| data, err := ioutil.ReadFile(outfile) |
| if err != nil { |
| t.Fatalf("can't read output file %s. %q", outfile, err) |
| } |
| |
| if string(data) != arg { |
| t.Fatalf("bad data. Expected %q, got %q", data, arg) |
| } |
| } |
| |
| // getPackage returns the name of the current package, which makes running this |
| // test in a fork simpler |
| func getPackage() []byte { |
| pc, _, _, _ := runtime.Caller(0) |
| fullFuncName := runtime.FuncForPC(pc).Name() |
| idx := strings.LastIndex(fullFuncName, ".") |
| return []byte(fullFuncName[:idx]) // trim off function details |
| } |
| |
| var testprogleader = []byte(` |
| // Test program for atexit, gets output file and data as arguments and writes |
| // data to output file in atexit handler. |
| package main |
| |
| import ( |
| "`) |
| var testprogtrailer = []byte( |
| `" |
| "flag" |
| "fmt" |
| "io/ioutil" |
| ) |
| |
| var outfile = "" |
| var data = "" |
| |
| func handler() { |
| ioutil.WriteFile(outfile, []byte(data), 0666) |
| } |
| |
| func badHandler() { |
| n := 0 |
| fmt.Println(1/n) |
| } |
| |
| func main() { |
| flag.Parse() |
| outfile = flag.Arg(0) |
| data = flag.Arg(1) |
| |
| logrus.RegisterExitHandler(handler) |
| logrus.RegisterExitHandler(badHandler) |
| logrus.Fatal("Bye bye") |
| } |
| `) |