-
Notifications
You must be signed in to change notification settings - Fork 2
/
shutdown_test.go
148 lines (118 loc) · 4.14 KB
/
shutdown_test.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
140
141
142
143
144
145
146
147
148
package main
import (
"context"
"fmt"
"io"
"os"
"testing"
"time"
logrus "github.com/sirupsen/logrus"
"github.com/stretchr/testify/mock"
)
// MockGRPCServer is a mock implementation of the GRPCServer interface.
type MockGRPCServer struct {
mock.Mock
}
func (m *MockGRPCServer) GracefulStop() {
m.Called()
}
// MockHTTPServer is a mock implementation of the HTTPServer interface.
type MockHTTPServer struct {
mock.Mock
}
// Shutdown simulates the shutdown of an HTTP server.
//
// This method mocks the behavior of the HTTPServer's Shutdown method. It takes
// a context as an argument and returns an error, which can be predefined in
// the test case using testify's mock functionalities.
func (m *MockHTTPServer) Shutdown(ctx context.Context) error {
args := m.Called(ctx)
return args.Error(0)
}
// TestGracefulShutdown tests the gracefulShutdown function.
func TestGracefulShutdown(t *testing.T) {
// Suppress log output by setting the output to io.Discard.
logrus.SetOutput(io.Discard)
// Create mock servers.
mockGRPCServer := new(MockGRPCServer)
mockHTTPServer := new(MockHTTPServer)
mockPProfServer := new(MockHTTPServer)
// Setup expectations for the mock servers.
mockGRPCServer.On("GracefulStop").Return()
mockHTTPServer.On("Shutdown", mock.Anything).Return(nil)
mockPProfServer.On("Shutdown", mock.Anything).Return(nil)
// Create a signal channel.
sigChan := make(chan os.Signal, 1)
// Run gracefulShutdown in a separate goroutine.
go gracefulShutdown(
sigChan, mockGRPCServer, mockHTTPServer, mockPProfServer,
)
// Simulate sending an interrupt signal.
sigChan <- os.Interrupt
// Wait a bit to let the gracefulShutdown function complete.
time.Sleep(100 * time.Millisecond)
// Assert that all expectations were met.
mockGRPCServer.AssertExpectations(t)
mockHTTPServer.AssertExpectations(t)
mockPProfServer.AssertExpectations(t)
}
// TestGracefulShutdownHTTPError tests gracefulShutdown when the HTTP server
// returns an error during shutdown.
func TestGracefulShutdownHTTPError(t *testing.T) {
// Suppress log output by setting the output to io.Discard.
logrus.SetOutput(io.Discard)
// Create mock servers.
mockGRPCServer := new(MockGRPCServer)
mockHTTPServer := new(MockHTTPServer)
mockPProfServer := new(MockHTTPServer)
// Setup expectations for the mock servers.
mockGRPCServer.On("GracefulStop").Return()
mockHTTPServer.On("Shutdown", mock.Anything).Return(
fmt.Errorf("HTTP shutdown error"),
)
mockPProfServer.On("Shutdown", mock.Anything).Return(nil)
// Create a signal channel.
sigChan := make(chan os.Signal, 1)
// Run gracefulShutdown in a separate goroutine.
go gracefulShutdown(
sigChan, mockGRPCServer, mockHTTPServer, mockPProfServer,
)
// Simulate sending an interrupt signal.
sigChan <- os.Interrupt
// Wait a bit to let the gracefulShutdown function complete.
time.Sleep(100 * time.Millisecond)
// Assert that all expectations were met.
mockGRPCServer.AssertExpectations(t)
mockHTTPServer.AssertExpectations(t)
mockPProfServer.AssertExpectations(t)
}
// TestGracefulShutdownPProfError tests gracefulShutdown when the pprof server
// returns an error during shutdown.
func TestGracefulShutdownPProfError(t *testing.T) {
// Suppress log output by setting the output to io.Discard.
logrus.SetOutput(io.Discard)
// Create mock servers.
mockGRPCServer := new(MockGRPCServer)
mockHTTPServer := new(MockHTTPServer)
mockPProfServer := new(MockHTTPServer)
// Setup expectations for the mock servers.
mockGRPCServer.On("GracefulStop").Return()
mockHTTPServer.On("Shutdown", mock.Anything).Return(nil)
mockPProfServer.On("Shutdown", mock.Anything).Return(
fmt.Errorf("PProf shutdown error"),
)
// Create a signal channel.
sigChan := make(chan os.Signal, 1)
// Run gracefulShutdown in a separate goroutine.
go gracefulShutdown(
sigChan, mockGRPCServer, mockHTTPServer, mockPProfServer,
)
// Simulate sending an interrupt signal.
sigChan <- os.Interrupt
// Wait a bit to let the gracefulShutdown function complete.
time.Sleep(100 * time.Millisecond)
// Assert that all expectations were met.
mockGRPCServer.AssertExpectations(t)
mockHTTPServer.AssertExpectations(t)
mockPProfServer.AssertExpectations(t)
}