-
-
Notifications
You must be signed in to change notification settings - Fork 274
/
testing.go
139 lines (119 loc) Β· 3.28 KB
/
testing.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
// Package pmtesting provides a simple unit test setup routine.
//
// Usage:
//
// package name
//
// import (
// "testing"
//
// "github.com/safing/portmaster/core/pmtesting"
// )
//
// func TestMain(m *testing.M) {
// pmtesting.TestMain(m, module)
// }
//
package pmtesting
import (
"flag"
"fmt"
"os"
"path/filepath"
"runtime/pprof"
"testing"
"github.com/safing/portbase/dataroot"
"github.com/safing/portbase/log"
"github.com/safing/portbase/modules"
"github.com/safing/portmaster/core/base"
// module dependencies
_ "github.com/safing/portbase/database/storage/hashmap"
)
var (
printStackOnExit bool
)
func init() {
flag.BoolVar(&printStackOnExit, "print-stack-on-exit", false, "prints the stack before of shutting down")
}
// TestHookFunc describes the functions passed to TestMainWithHooks.
type TestHookFunc func() error
// TestMain provides a simple unit test setup routine.
func TestMain(m *testing.M, module *modules.Module) {
TestMainWithHooks(m, module, nil, nil)
}
// TestMainWithHooks provides a simple unit test setup routine and calls
// afterStartFn after modules have started and beforeStopFn before modules
// are shutdown.
func TestMainWithHooks(m *testing.M, module *modules.Module, afterStartFn, beforeStopFn TestHookFunc) {
// enable module for testing
if module != nil {
module.Enable()
}
// switch databases to memory only
base.DefaultDatabaseStorageType = "hashmap"
// switch API to high port
base.DefaultAPIListenAddress = "127.0.0.1:10817"
// set log level
log.SetLogLevel(log.TraceLevel)
// tmp dir for data root (db & config)
tmpDir := filepath.Join(os.TempDir(), "portmaster-testing")
// initialize data dir
err := dataroot.Initialize(tmpDir, 0755)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to initialize data root: %s\n", err)
os.Exit(1)
}
// start modules
var exitCode int
err = modules.Start()
if err != nil {
// starting failed
fmt.Fprintf(os.Stderr, "failed to setup test: %s\n", err)
exitCode = 1
} else {
runTests := true
if afterStartFn != nil {
if err := afterStartFn(); err != nil {
fmt.Fprintf(os.Stderr, "failed to run test start hook: %s\n", err)
runTests = false
exitCode = 1
}
}
if runTests {
// run tests
exitCode = m.Run()
}
}
if beforeStopFn != nil {
if err := beforeStopFn(); err != nil {
fmt.Fprintf(os.Stderr, "failed to run test shutdown hook: %s\n", err)
}
}
// shutdown
_ = modules.Shutdown()
if modules.GetExitStatusCode() != 0 {
exitCode = modules.GetExitStatusCode()
fmt.Fprintf(os.Stderr, "failed to cleanly shutdown test: %s\n", err)
}
printStack()
// clean up and exit
// Important: Do not remove tmpDir, as it is used as a cache for updates.
// remove config
_ = os.Remove(filepath.Join(tmpDir, "config.json"))
// remove databases
_ = os.Remove(filepath.Join(tmpDir, "databases.json"))
_ = os.RemoveAll(filepath.Join(tmpDir, "databases"))
os.Exit(exitCode)
}
func printStack() {
if printStackOnExit {
fmt.Println("=== PRINTING TRACES ===")
fmt.Println("=== GOROUTINES ===")
_ = pprof.Lookup("goroutine").WriteTo(os.Stdout, 2)
fmt.Println("=== BLOCKING ===")
_ = pprof.Lookup("block").WriteTo(os.Stdout, 2)
fmt.Println("=== MUTEXES ===")
_ = pprof.Lookup("mutex").WriteTo(os.Stdout, 2)
fmt.Println("=== END TRACES ===")
}
}