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