1package logrus 2 3import ( 4 "io/ioutil" 5 "log" 6 "os" 7 "os/exec" 8 "path/filepath" 9 "runtime" 10 "strings" 11 "testing" 12 "time" 13) 14 15func TestRegister(t *testing.T) { 16 current := len(handlers) 17 18 var results []string 19 20 h1 := func() { results = append(results, "first") } 21 h2 := func() { results = append(results, "second") } 22 23 RegisterExitHandler(h1) 24 RegisterExitHandler(h2) 25 26 if len(handlers) != current+2 { 27 t.Fatalf("expected %d handlers, got %d", current+2, len(handlers)) 28 } 29 30 runHandlers() 31 32 if len(results) != 2 { 33 t.Fatalf("expected 2 handlers to be run, ran %d", len(results)) 34 } 35 36 if results[0] != "first" { 37 t.Fatal("expected handler h1 to be run first, but it wasn't") 38 } 39 40 if results[1] != "second" { 41 t.Fatal("expected handler h2 to be run second, but it wasn't") 42 } 43} 44 45func TestDefer(t *testing.T) { 46 current := len(handlers) 47 48 var results []string 49 50 h1 := func() { results = append(results, "first") } 51 h2 := func() { results = append(results, "second") } 52 53 DeferExitHandler(h1) 54 DeferExitHandler(h2) 55 56 if len(handlers) != current+2 { 57 t.Fatalf("expected %d handlers, got %d", current+2, len(handlers)) 58 } 59 60 runHandlers() 61 62 if len(results) != 2 { 63 t.Fatalf("expected 2 handlers to be run, ran %d", len(results)) 64 } 65 66 if results[0] != "second" { 67 t.Fatal("expected handler h2 to be run first, but it wasn't") 68 } 69 70 if results[1] != "first" { 71 t.Fatal("expected handler h1 to be run second, but it wasn't") 72 } 73} 74 75func TestHandler(t *testing.T) { 76 testprog := testprogleader 77 testprog = append(testprog, getPackage()...) 78 testprog = append(testprog, testprogtrailer...) 79 tempDir, err := ioutil.TempDir("", "test_handler") 80 if err != nil { 81 log.Fatalf("can't create temp dir. %q", err) 82 } 83 defer os.RemoveAll(tempDir) 84 85 gofile := filepath.Join(tempDir, "gofile.go") 86 if err := ioutil.WriteFile(gofile, testprog, 0666); err != nil { 87 t.Fatalf("can't create go file. %q", err) 88 } 89 90 outfile := filepath.Join(tempDir, "outfile.out") 91 arg := time.Now().UTC().String() 92 err = exec.Command("go", "run", gofile, outfile, arg).Run() 93 if err == nil { 94 t.Fatalf("completed normally, should have failed") 95 } 96 97 data, err := ioutil.ReadFile(outfile) 98 if err != nil { 99 t.Fatalf("can't read output file %s. %q", outfile, err) 100 } 101 102 if string(data) != arg { 103 t.Fatalf("bad data. Expected %q, got %q", data, arg) 104 } 105} 106 107// getPackage returns the name of the current package, which makes running this 108// test in a fork simpler 109func getPackage() []byte { 110 pc, _, _, _ := runtime.Caller(0) 111 fullFuncName := runtime.FuncForPC(pc).Name() 112 idx := strings.LastIndex(fullFuncName, ".") 113 return []byte(fullFuncName[:idx]) // trim off function details 114} 115 116var testprogleader = []byte(` 117// Test program for atexit, gets output file and data as arguments and writes 118// data to output file in atexit handler. 119package main 120 121import ( 122 "`) 123var testprogtrailer = []byte( 124 `" 125 "flag" 126 "fmt" 127 "io/ioutil" 128) 129 130var outfile = "" 131var data = "" 132 133func handler() { 134 ioutil.WriteFile(outfile, []byte(data), 0666) 135} 136 137func badHandler() { 138 n := 0 139 fmt.Println(1/n) 140} 141 142func main() { 143 flag.Parse() 144 outfile = flag.Arg(0) 145 data = flag.Arg(1) 146 147 logrus.RegisterExitHandler(handler) 148 logrus.RegisterExitHandler(badHandler) 149 logrus.Fatal("Bye bye") 150} 151`) 152