1package logrus
2
3import (
4	"io/ioutil"
5	"log"
6	"os"
7	"os/exec"
8	"path/filepath"
9	"testing"
10	"time"
11)
12
13func TestRegister(t *testing.T) {
14	current := len(handlers)
15	RegisterExitHandler(func() {})
16	if len(handlers) != current+1 {
17		t.Fatalf("expected %d handlers, got %d", current+1, len(handlers))
18	}
19}
20
21func TestHandler(t *testing.T) {
22	tempDir, err := ioutil.TempDir("", "test_handler")
23	if err != nil {
24		log.Fatalf("can't create temp dir. %q", err)
25	}
26	defer os.RemoveAll(tempDir)
27
28	gofile := filepath.Join(tempDir, "gofile.go")
29	if err := ioutil.WriteFile(gofile, testprog, 0666); err != nil {
30		t.Fatalf("can't create go file. %q", err)
31	}
32
33	outfile := filepath.Join(tempDir, "outfile.out")
34	arg := time.Now().UTC().String()
35	err = exec.Command("go", "run", gofile, outfile, arg).Run()
36	if err == nil {
37		t.Fatalf("completed normally, should have failed")
38	}
39
40	data, err := ioutil.ReadFile(outfile)
41	if err != nil {
42		t.Fatalf("can't read output file %s. %q", outfile, err)
43	}
44
45	if string(data) != arg {
46		t.Fatalf("bad data. Expected %q, got %q", data, arg)
47	}
48}
49
50var testprog = []byte(`
51// Test program for atexit, gets output file and data as arguments and writes
52// data to output file in atexit handler.
53package main
54
55import (
56	"github.com/sirupsen/logrus"
57	"flag"
58	"fmt"
59	"io/ioutil"
60)
61
62var outfile = ""
63var data = ""
64
65func handler() {
66	ioutil.WriteFile(outfile, []byte(data), 0666)
67}
68
69func badHandler() {
70	n := 0
71	fmt.Println(1/n)
72}
73
74func main() {
75	flag.Parse()
76	outfile = flag.Arg(0)
77	data = flag.Arg(1)
78
79	logrus.RegisterExitHandler(handler)
80	logrus.RegisterExitHandler(badHandler)
81	logrus.Fatal("Bye bye")
82}
83`)
84