kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package agentserver
15

16
import (
17
	"bytes"
18
	"context"
19
	"encoding/json"
20
	"errors"
21
	"fmt"
22
	"io/ioutil"
23
	"net/url"
24
	"testing"
25
	"time"
26

27
	"github.com/uber/kraken/agent/agentclient"
28
	"github.com/uber/kraken/build-index/tagclient"
29
	"github.com/uber/kraken/core"
30
	"github.com/uber/kraken/lib/store"
31
	"github.com/uber/kraken/lib/torrent/scheduler"
32
	"github.com/uber/kraken/lib/torrent/scheduler/connstate"
33
	mocktagclient "github.com/uber/kraken/mocks/build-index/tagclient"
34
	mockcontainerruntime "github.com/uber/kraken/mocks/lib/containerruntime"
35
	mockcontainerd "github.com/uber/kraken/mocks/lib/containerruntime/containerd"
36
	mockdockerdaemon "github.com/uber/kraken/mocks/lib/containerruntime/dockerdaemon"
37
	mockscheduler "github.com/uber/kraken/mocks/lib/torrent/scheduler"
38
	"github.com/uber/kraken/utils/httputil"
39
	"github.com/uber/kraken/utils/testutil"
40

41
	"github.com/golang/mock/gomock"
42
	"github.com/stretchr/testify/require"
43
	"github.com/uber-go/tally"
44
)
45

46
type serverMocks struct {
47
	cads             *store.CADownloadStore
48
	sched            *mockscheduler.MockReloadableScheduler
49
	tags             *mocktagclient.MockClient
50
	dockerCli        *mockdockerdaemon.MockDockerClient
51
	containerdCli    *mockcontainerd.MockClient
52
	containerRuntime *mockcontainerruntime.MockFactory
53
	cleanup          *testutil.Cleanup
54
}
55

56
func newServerMocks(t *testing.T) (*serverMocks, func()) {
57
	var cleanup testutil.Cleanup
58

59
	cads, c := store.CADownloadStoreFixture()
60
	cleanup.Add(c)
61

62
	ctrl := gomock.NewController(t)
63
	cleanup.Add(ctrl.Finish)
64

65
	sched := mockscheduler.NewMockReloadableScheduler(ctrl)
66

67
	tags := mocktagclient.NewMockClient(ctrl)
68

69
	dockerCli := mockdockerdaemon.NewMockDockerClient(ctrl)
70
	containerdCli := mockcontainerd.NewMockClient(ctrl)
71
	containerruntime := mockcontainerruntime.NewMockFactory(ctrl)
72
	return &serverMocks{
73
		cads, sched, tags, dockerCli, containerdCli,
74
		containerruntime, &cleanup}, cleanup.Run
75
}
76

77
func (m *serverMocks) startServer() string {
78
	s := New(Config{}, tally.NoopScope, m.cads, m.sched, m.tags, m.containerRuntime)
79
	addr, stop := testutil.StartServer(s.Handler())
80
	m.cleanup.Add(stop)
81
	return addr
82
}
83

84
func TestGetTag(t *testing.T) {
85
	require := require.New(t)
86

87
	mocks, cleanup := newServerMocks(t)
88
	defer cleanup()
89

90
	tag := core.TagFixture()
91
	d := core.DigestFixture()
92

93
	mocks.tags.EXPECT().Get(tag).Return(d, nil)
94

95
	c := agentclient.New(mocks.startServer())
96

97
	result, err := c.GetTag(tag)
98
	require.NoError(err)
99
	require.Equal(d, result)
100
}
101

102
func TestGetTagNotFound(t *testing.T) {
103
	require := require.New(t)
104

105
	mocks, cleanup := newServerMocks(t)
106
	defer cleanup()
107

108
	tag := core.TagFixture()
109

110
	mocks.tags.EXPECT().Get(tag).Return(core.Digest{}, tagclient.ErrTagNotFound)
111

112
	c := agentclient.New(mocks.startServer())
113

114
	_, err := c.GetTag(tag)
115
	require.Error(err)
116
	require.Equal(agentclient.ErrTagNotFound, err)
117
}
118

119
func TestDownload(t *testing.T) {
120
	require := require.New(t)
121

122
	mocks, cleanup := newServerMocks(t)
123
	defer cleanup()
124

125
	namespace := core.TagFixture()
126
	blob := core.NewBlobFixture()
127

128
	mocks.sched.EXPECT().Download(namespace, blob.Digest).DoAndReturn(
129
		func(namespace string, d core.Digest) error {
130
			return store.RunDownload(mocks.cads, d, blob.Content)
131
		})
132

133
	addr := mocks.startServer()
134
	c := agentclient.New(addr)
135

136
	r, err := c.Download(namespace, blob.Digest)
137
	require.NoError(err)
138
	result, err := ioutil.ReadAll(r)
139
	require.NoError(err)
140
	require.Equal(string(blob.Content), string(result))
141
}
142

143
func TestDownloadNotFound(t *testing.T) {
144
	require := require.New(t)
145

146
	mocks, cleanup := newServerMocks(t)
147
	defer cleanup()
148

149
	namespace := core.TagFixture()
150
	blob := core.NewBlobFixture()
151

152
	mocks.sched.EXPECT().Download(namespace, blob.Digest).Return(scheduler.ErrTorrentNotFound)
153

154
	addr := mocks.startServer()
155
	c := agentclient.New(addr)
156

157
	_, err := c.Download(namespace, blob.Digest)
158
	require.Error(err)
159
	require.True(httputil.IsNotFound(err))
160
}
161

162
func TestDownloadUnknownError(t *testing.T) {
163
	require := require.New(t)
164

165
	mocks, cleanup := newServerMocks(t)
166
	defer cleanup()
167

168
	namespace := core.TagFixture()
169
	blob := core.NewBlobFixture()
170

171
	mocks.sched.EXPECT().Download(namespace, blob.Digest).Return(fmt.Errorf("test error"))
172

173
	addr := mocks.startServer()
174
	c := agentclient.New(addr)
175

176
	_, err := c.Download(namespace, blob.Digest)
177
	require.Error(err)
178
	require.True(httputil.IsStatus(err, 500))
179
}
180

181
func TestHealthHandler(t *testing.T) {
182
	tests := []struct {
183
		desc     string
184
		probeErr error
185
	}{
186
		{"probe error", errors.New("some probe error")},
187
		{"healthy", nil},
188
	}
189
	for _, test := range tests {
190
		t.Run(test.desc, func(t *testing.T) {
191
			require := require.New(t)
192

193
			mocks, cleanup := newServerMocks(t)
194
			defer cleanup()
195

196
			mocks.sched.EXPECT().Probe().Return(test.probeErr)
197

198
			addr := mocks.startServer()
199

200
			_, err := httputil.Get(fmt.Sprintf("http://%s/health", addr))
201
			if test.probeErr != nil {
202
				require.Error(err)
203
			} else {
204
				require.NoError(err)
205
			}
206
		})
207
	}
208
}
209

210
func TestPatchSchedulerConfigHandler(t *testing.T) {
211
	require := require.New(t)
212

213
	mocks, cleanup := newServerMocks(t)
214
	defer cleanup()
215

216
	addr := mocks.startServer()
217

218
	config := scheduler.Config{
219
		ConnTTI: time.Minute,
220
	}
221
	b, err := json.Marshal(config)
222
	require.NoError(err)
223

224
	mocks.sched.EXPECT().Reload(config)
225

226
	_, err = httputil.Patch(
227
		fmt.Sprintf("http://%s/x/config/scheduler", addr),
228
		httputil.SendBody(bytes.NewReader(b)))
229
	require.NoError(err)
230
}
231

232
func TestGetBlacklistHandler(t *testing.T) {
233
	require := require.New(t)
234

235
	mocks, cleanup := newServerMocks(t)
236
	defer cleanup()
237

238
	blacklist := []connstate.BlacklistedConn{{
239
		PeerID:    core.PeerIDFixture(),
240
		InfoHash:  core.InfoHashFixture(),
241
		Remaining: time.Second,
242
	}}
243
	mocks.sched.EXPECT().BlacklistSnapshot().Return(blacklist, nil)
244

245
	addr := mocks.startServer()
246

247
	resp, err := httputil.Get(fmt.Sprintf("http://%s/x/blacklist", addr))
248
	require.NoError(err)
249

250
	var result []connstate.BlacklistedConn
251
	require.NoError(json.NewDecoder(resp.Body).Decode(&result))
252
	require.Equal(blacklist, result)
253
}
254

255
func TestDeleteBlobHandler(t *testing.T) {
256
	require := require.New(t)
257

258
	mocks, cleanup := newServerMocks(t)
259
	defer cleanup()
260

261
	d := core.DigestFixture()
262

263
	addr := mocks.startServer()
264

265
	mocks.sched.EXPECT().RemoveTorrent(d).Return(nil)
266

267
	_, err := httputil.Delete(fmt.Sprintf("http://%s/blobs/%s", addr, d))
268
	require.NoError(err)
269
}
270

271
func TestPreloadHandler(t *testing.T) {
272
	tag := url.PathEscape("repo1:tag1")
273
	tests := []struct {
274
		name          string
275
		url           string
276
		setup         func(*serverMocks)
277
		expectedError string
278
	}{
279
		{
280
			name: "success docker",
281
			url:  fmt.Sprintf("/preload/tags/%s", tag),
282
			setup: func(mocks *serverMocks) {
283
				mocks.dockerCli.EXPECT().
284
					PullImage(context.Background(), "repo1", "tag1").Return(nil)
285
				mocks.containerRuntime.EXPECT().
286
					DockerClient().Return(mocks.dockerCli)
287
			},
288
		},
289
		{
290
			name: "success containerd",
291
			url:  fmt.Sprintf("/preload/tags/%s?runtime=containerd&namespace=name.space1", tag),
292
			setup: func(mocks *serverMocks) {
293
				mocks.containerdCli.EXPECT().
294
					PullImage(context.Background(), "name.space1", "repo1", "tag1").Return(nil)
295
				mocks.containerRuntime.EXPECT().
296
					ContainerdClient().Return(mocks.containerdCli)
297
			},
298
		},
299
		{
300
			name:          "unsupported runtime",
301
			url:           fmt.Sprintf("/preload/tags/%s?runtime=crio", tag),
302
			setup:         func(_ *serverMocks) {},
303
			expectedError: "/preload/tags/repo1:tag1?runtime=crio 500: unsupported container runtime",
304
		},
305
	}
306

307
	for _, tt := range tests {
308
		t.Run(tt.name, func(t *testing.T) {
309
			require := require.New(t)
310

311
			mocks, cleanup := newServerMocks(t)
312
			defer cleanup()
313

314
			tt.setup(mocks)
315
			addr := mocks.startServer()
316

317
			_, err := httputil.Get(fmt.Sprintf("http://%s%s", addr, tt.url))
318
			if tt.expectedError != "" {
319
				require.EqualError(err,
320
					fmt.Sprintf("GET http://%s%s", addr, tt.expectedError))
321
			} else {
322
				require.NoError(err)
323
			}
324
		})
325
	}
326
}
327