forked from gravitational/teleport
-
Notifications
You must be signed in to change notification settings - Fork 0
/
suite.go
179 lines (153 loc) · 5.38 KB
/
suite.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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
/*
* Teleport
* Copyright (C) 2023 Gravitational, Inc.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package integration
import (
"context"
"os"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
"github.com/gravitational/teleport/integrations/lib/logger"
)
// Suite is a basic testing suite enhanced with context management.
type Suite struct {
suite.Suite
contexts map[*testing.T]contexts
}
// AppI is an app that can be spawned along with running test.
type AppI interface {
// Run starts the application
Run(ctx context.Context) error
// WaitReady waits till the application finishes initialization
WaitReady(ctx context.Context) (bool, error)
// Err returns last error
Err() error
// Shutdown shuts the application down
Shutdown(ctx context.Context) error
}
type contexts struct {
// baseCtx is the the base context for appCtx and testCtx.
// It could store some test-specific information stored using context.WithValue()
// such as test name for the logger etc.
baseCtx context.Context
// appCtx inherits from baseCtx. Its purpose is to limit the lifetime of the apps running in parallel.
// By "app" we mean some plugin (e.g. access/slack) or the Teleport process (lib/testing/integration package).
// Its timeout is slightly higher than testCtx's for a reason. When the test example fails with timeout
// we want to see the exact line of the test file where the fail took place. But if the app dies at the same time
// as the some operation in the test example we probably we'll see the line where the app failed, not the test
// which is non-informative but we really want to see what line of the test caused the timeout and where it happened.
appCtx context.Context
// testCtx inherits from baseCtx. Its purpose is to limit the lifetime of the test method.
// This context is guaranteed to be canceled earlier than appCtx for better error reporting (see explanation above).
testCtx context.Context
}
// SetT sets the current *testing.T context.
func (s *Suite) SetT(t *testing.T) {
oldT := s.T()
s.Suite.SetT(t)
s.initContexts(oldT, t)
}
func (s *Suite) initContexts(oldT *testing.T, newT *testing.T) {
if s.contexts == nil {
s.contexts = make(map[*testing.T]contexts)
}
contexts, ok := s.contexts[newT]
if ok {
// Context already initialized.
// This happens when testify sets the parent context back after running a subtest.
return
}
var baseCtx context.Context
if oldT != nil && strings.HasPrefix(newT.Name(), oldT.Name()+"/") {
// We are running a subtest so lets inherit the context too.
baseCtx = s.contexts[oldT].testCtx
} else {
baseCtx = context.Background()
}
baseCtx, _ = logger.WithField(baseCtx, "test", newT.Name())
baseCtx, cancel := context.WithCancel(baseCtx)
newT.Cleanup(cancel)
contexts.baseCtx = baseCtx
contexts.appCtx = baseCtx
contexts.testCtx = baseCtx
// Just memoize the context in a map and that's all.
// Lets not bother with cleaning up this storage, it's not gonna be that big.
s.contexts[newT] = contexts
}
// SetContextTimeout limits the lifetime of test and app contexts.
func (s *Suite) SetContextTimeout(timeout time.Duration) context.Context {
t := s.T()
t.Helper()
contexts, ok := s.contexts[t]
require.True(t, ok)
var cancel context.CancelFunc
// We set appCtx timeout slightly higher than testCtx for test assertions to fall earlier than
// app (plugin) fails.
contexts.appCtx, cancel = context.WithTimeout(contexts.baseCtx, timeout+500*time.Millisecond)
t.Cleanup(cancel)
contexts.testCtx, cancel = context.WithTimeout(contexts.baseCtx, timeout)
t.Cleanup(cancel)
s.contexts[t] = contexts
return contexts.testCtx
}
// Context returns a current test context.
func (s *Suite) Context() context.Context {
t := s.T()
t.Helper()
contexts, ok := s.contexts[t]
require.True(t, ok)
return contexts.testCtx
}
// NewTmpFile creates a new temporary file.
func (s *Suite) NewTmpFile(pattern string) *os.File {
t := s.T()
t.Helper()
file, err := os.CreateTemp("", pattern)
require.NoError(t, err)
t.Cleanup(func() {
err := os.Remove(file.Name())
require.NoError(t, err)
})
return file
}
// StartApp spawns an app in parallel with the running test/suite.
func (s *Suite) StartApp(app AppI) {
t := s.T()
t.Helper()
contexts, ok := s.contexts[t]
require.True(t, ok)
go func() {
ctx := contexts.appCtx
if err := app.Run(ctx); err != nil {
// We're in a goroutine so we can't just require.NoError(t, err).
// All we can do is to log an error.
logger.Get(ctx).WithError(err).Error("Application failed")
}
}()
t.Cleanup(func() {
err := app.Shutdown(contexts.appCtx)
assert.NoError(t, err)
assert.NoError(t, app.Err())
})
ok, err := app.WaitReady(contexts.testCtx)
require.NoError(t, err)
require.True(t, ok)
}