github.com/uber/kraken@v0.1.4/agent/agentserver/server_test.go (about)

     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  	mockdockerdaemon "github.com/uber/kraken/mocks/lib/dockerdaemon"
    35  	mockscheduler "github.com/uber/kraken/mocks/lib/torrent/scheduler"
    36  	"github.com/uber/kraken/utils/httputil"
    37  	"github.com/uber/kraken/utils/testutil"
    38  
    39  	"github.com/golang/mock/gomock"
    40  	"github.com/stretchr/testify/require"
    41  	"github.com/uber-go/tally"
    42  )
    43  
    44  type serverMocks struct {
    45  	cads      *store.CADownloadStore
    46  	sched     *mockscheduler.MockReloadableScheduler
    47  	tags      *mocktagclient.MockClient
    48  	dockerCli *mockdockerdaemon.MockDockerClient
    49  	cleanup   *testutil.Cleanup
    50  }
    51  
    52  func newServerMocks(t *testing.T) (*serverMocks, func()) {
    53  	var cleanup testutil.Cleanup
    54  
    55  	cads, c := store.CADownloadStoreFixture()
    56  	cleanup.Add(c)
    57  
    58  	ctrl := gomock.NewController(t)
    59  	cleanup.Add(ctrl.Finish)
    60  
    61  	sched := mockscheduler.NewMockReloadableScheduler(ctrl)
    62  
    63  	tags := mocktagclient.NewMockClient(ctrl)
    64  
    65  	dockerCli := mockdockerdaemon.NewMockDockerClient(ctrl)
    66  
    67  	return &serverMocks{cads, sched, tags, dockerCli, &cleanup}, cleanup.Run
    68  }
    69  
    70  func (m *serverMocks) startServer() string {
    71  	s := New(Config{}, tally.NoopScope, m.cads, m.sched, m.tags, m.dockerCli)
    72  	addr, stop := testutil.StartServer(s.Handler())
    73  	m.cleanup.Add(stop)
    74  	return addr
    75  }
    76  
    77  func TestGetTag(t *testing.T) {
    78  	require := require.New(t)
    79  
    80  	mocks, cleanup := newServerMocks(t)
    81  	defer cleanup()
    82  
    83  	tag := core.TagFixture()
    84  	d := core.DigestFixture()
    85  
    86  	mocks.tags.EXPECT().Get(tag).Return(d, nil)
    87  
    88  	c := agentclient.New(mocks.startServer())
    89  
    90  	result, err := c.GetTag(tag)
    91  	require.NoError(err)
    92  	require.Equal(d, result)
    93  }
    94  
    95  func TestGetTagNotFound(t *testing.T) {
    96  	require := require.New(t)
    97  
    98  	mocks, cleanup := newServerMocks(t)
    99  	defer cleanup()
   100  
   101  	tag := core.TagFixture()
   102  
   103  	mocks.tags.EXPECT().Get(tag).Return(core.Digest{}, tagclient.ErrTagNotFound)
   104  
   105  	c := agentclient.New(mocks.startServer())
   106  
   107  	_, err := c.GetTag(tag)
   108  	require.Error(err)
   109  	require.Equal(agentclient.ErrTagNotFound, err)
   110  }
   111  
   112  func TestDownload(t *testing.T) {
   113  	require := require.New(t)
   114  
   115  	mocks, cleanup := newServerMocks(t)
   116  	defer cleanup()
   117  
   118  	namespace := core.TagFixture()
   119  	blob := core.NewBlobFixture()
   120  
   121  	mocks.sched.EXPECT().Download(namespace, blob.Digest).DoAndReturn(
   122  		func(namespace string, d core.Digest) error {
   123  			return store.RunDownload(mocks.cads, d, blob.Content)
   124  		})
   125  
   126  	addr := mocks.startServer()
   127  	c := agentclient.New(addr)
   128  
   129  	r, err := c.Download(namespace, blob.Digest)
   130  	require.NoError(err)
   131  	result, err := ioutil.ReadAll(r)
   132  	require.NoError(err)
   133  	require.Equal(string(blob.Content), string(result))
   134  }
   135  
   136  func TestDownloadNotFound(t *testing.T) {
   137  	require := require.New(t)
   138  
   139  	mocks, cleanup := newServerMocks(t)
   140  	defer cleanup()
   141  
   142  	namespace := core.TagFixture()
   143  	blob := core.NewBlobFixture()
   144  
   145  	mocks.sched.EXPECT().Download(namespace, blob.Digest).Return(scheduler.ErrTorrentNotFound)
   146  
   147  	addr := mocks.startServer()
   148  	c := agentclient.New(addr)
   149  
   150  	_, err := c.Download(namespace, blob.Digest)
   151  	require.Error(err)
   152  	require.True(httputil.IsNotFound(err))
   153  }
   154  
   155  func TestDownloadUnknownError(t *testing.T) {
   156  	require := require.New(t)
   157  
   158  	mocks, cleanup := newServerMocks(t)
   159  	defer cleanup()
   160  
   161  	namespace := core.TagFixture()
   162  	blob := core.NewBlobFixture()
   163  
   164  	mocks.sched.EXPECT().Download(namespace, blob.Digest).Return(fmt.Errorf("test error"))
   165  
   166  	addr := mocks.startServer()
   167  	c := agentclient.New(addr)
   168  
   169  	_, err := c.Download(namespace, blob.Digest)
   170  	require.Error(err)
   171  	require.True(httputil.IsStatus(err, 500))
   172  }
   173  
   174  func TestHealthHandler(t *testing.T) {
   175  	tests := []struct {
   176  		desc     string
   177  		probeErr error
   178  	}{
   179  		{"probe error", errors.New("some probe error")},
   180  		{"healthy", nil},
   181  	}
   182  	for _, test := range tests {
   183  		t.Run(test.desc, func(t *testing.T) {
   184  			require := require.New(t)
   185  
   186  			mocks, cleanup := newServerMocks(t)
   187  			defer cleanup()
   188  
   189  			mocks.sched.EXPECT().Probe().Return(test.probeErr)
   190  
   191  			addr := mocks.startServer()
   192  
   193  			_, err := httputil.Get(fmt.Sprintf("http://%s/health", addr))
   194  			if test.probeErr != nil {
   195  				require.Error(err)
   196  			} else {
   197  				require.NoError(err)
   198  			}
   199  		})
   200  	}
   201  }
   202  
   203  func TestPatchSchedulerConfigHandler(t *testing.T) {
   204  	require := require.New(t)
   205  
   206  	mocks, cleanup := newServerMocks(t)
   207  	defer cleanup()
   208  
   209  	addr := mocks.startServer()
   210  
   211  	config := scheduler.Config{
   212  		ConnTTI: time.Minute,
   213  	}
   214  	b, err := json.Marshal(config)
   215  	require.NoError(err)
   216  
   217  	mocks.sched.EXPECT().Reload(config)
   218  
   219  	_, err = httputil.Patch(
   220  		fmt.Sprintf("http://%s/x/config/scheduler", addr),
   221  		httputil.SendBody(bytes.NewReader(b)))
   222  	require.NoError(err)
   223  }
   224  
   225  func TestGetBlacklistHandler(t *testing.T) {
   226  	require := require.New(t)
   227  
   228  	mocks, cleanup := newServerMocks(t)
   229  	defer cleanup()
   230  
   231  	blacklist := []connstate.BlacklistedConn{{
   232  		PeerID:    core.PeerIDFixture(),
   233  		InfoHash:  core.InfoHashFixture(),
   234  		Remaining: time.Second,
   235  	}}
   236  	mocks.sched.EXPECT().BlacklistSnapshot().Return(blacklist, nil)
   237  
   238  	addr := mocks.startServer()
   239  
   240  	resp, err := httputil.Get(fmt.Sprintf("http://%s/x/blacklist", addr))
   241  	require.NoError(err)
   242  
   243  	var result []connstate.BlacklistedConn
   244  	require.NoError(json.NewDecoder(resp.Body).Decode(&result))
   245  	require.Equal(blacklist, result)
   246  }
   247  
   248  func TestDeleteBlobHandler(t *testing.T) {
   249  	require := require.New(t)
   250  
   251  	mocks, cleanup := newServerMocks(t)
   252  	defer cleanup()
   253  
   254  	d := core.DigestFixture()
   255  
   256  	addr := mocks.startServer()
   257  
   258  	mocks.sched.EXPECT().RemoveTorrent(d).Return(nil)
   259  
   260  	_, err := httputil.Delete(fmt.Sprintf("http://%s/blobs/%s", addr, d))
   261  	require.NoError(err)
   262  }
   263  
   264  func TestPreloadHandler(t *testing.T) {
   265  	require := require.New(t)
   266  
   267  	mocks, cleanup := newServerMocks(t)
   268  	defer cleanup()
   269  
   270  	addr := mocks.startServer()
   271  
   272  	tag := url.PathEscape("repo1:tag1")
   273  
   274  	mocks.dockerCli.EXPECT().PullImage(context.Background(), "repo1", "tag1").Return(nil)
   275  
   276  	_, err := httputil.Get(fmt.Sprintf("http://%s/preload/tags/%s", addr, tag))
   277  	require.NoError(err)
   278  }