github.com/uber/kraken@v0.1.4/lib/backend/manager_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 backend_test
    15  
    16  import (
    17  	"testing"
    18  
    19  	. "github.com/uber/kraken/lib/backend"
    20  	"github.com/uber/kraken/lib/backend/namepath"
    21  	"github.com/uber/kraken/lib/backend/testfs"
    22  	"github.com/uber/kraken/mocks/lib/backend"
    23  	"github.com/uber/kraken/utils/bandwidth"
    24  	"github.com/uber/kraken/utils/stringset"
    25  
    26  	"github.com/golang/mock/gomock"
    27  	"github.com/stretchr/testify/require"
    28  	"gopkg.in/yaml.v2"
    29  )
    30  
    31  func TestManagerNamespaceMatching(t *testing.T) {
    32  	ctrl := gomock.NewController(t)
    33  	defer ctrl.Finish()
    34  
    35  	c1 := mockbackend.NewMockClient(ctrl)
    36  	c2 := mockbackend.NewMockClient(ctrl)
    37  
    38  	tests := []struct {
    39  		namespace string
    40  		expected  Client
    41  	}{
    42  		{"static", c1},
    43  		{"namespace-foo/repo-bar", c2},
    44  	}
    45  	for _, test := range tests {
    46  		t.Run(test.namespace, func(t *testing.T) {
    47  			require := require.New(t)
    48  
    49  			m := ManagerFixture()
    50  
    51  			require.NoError(m.Register("static", c1))
    52  			require.NoError(m.Register("namespace-foo/.*", c2))
    53  
    54  			result, err := m.GetClient(test.namespace)
    55  			require.NoError(err)
    56  			require.True(
    57  				test.expected.(*mockbackend.MockClient) == result.(*mockbackend.MockClient))
    58  		})
    59  	}
    60  }
    61  
    62  func TestManagerErrDuplicateNamespace(t *testing.T) {
    63  	require := require.New(t)
    64  
    65  	c := &NoopClient{}
    66  	m := ManagerFixture()
    67  	require.NoError(m.Register("static", c))
    68  	require.Error(m.Register("static", c))
    69  }
    70  
    71  func TestManagerErrNamespaceNotFound(t *testing.T) {
    72  	require := require.New(t)
    73  
    74  	m := ManagerFixture()
    75  	_, err := m.GetClient("no-match")
    76  	require.Equal(ErrNamespaceNotFound, err)
    77  }
    78  
    79  func TestManagerNamespaceOrdering(t *testing.T) {
    80  	require := require.New(t)
    81  
    82  	fooBarAddr := "testfs-foo-bar"
    83  	fooAddr := "testfs-foo"
    84  	defaultAddr := "testfs-default"
    85  
    86  	configStr := `
    87  - namespace: foo/bar/.*
    88    backend:
    89        testfs:
    90            addr: testfs-foo-bar
    91            name_path: identity
    92  - namespace: foo/.*
    93    backend:
    94        testfs:
    95            addr: testfs-foo
    96            name_path: identity
    97  - namespace: .*
    98    backend:
    99        testfs:
   100            addr: testfs-default
   101            name_path: identity
   102  `
   103  	var configs []Config
   104  	require.NoError(yaml.Unmarshal([]byte(configStr), &configs))
   105  
   106  	m, err := NewManager(configs, AuthConfig{})
   107  	require.NoError(err)
   108  
   109  	for ns, expected := range map[string]string{
   110  		"foo/bar/baz": fooBarAddr,
   111  		"foo/bar/123": fooBarAddr,
   112  		"foo/123":     fooAddr,
   113  		"abc":         defaultAddr,
   114  		"xyz":         defaultAddr,
   115  		"bar/baz":     defaultAddr,
   116  		"":            defaultAddr,
   117  	} {
   118  		c, err := m.GetClient(ns)
   119  		require.NoError(err)
   120  		require.Equal(expected, c.(*testfs.Client).Addr(), "Namespace: %s", ns)
   121  	}
   122  }
   123  
   124  func TestManagerBandwidth(t *testing.T) {
   125  	require := require.New(t)
   126  
   127  	m, err := NewManager([]Config{{
   128  		Namespace: ".*",
   129  		Bandwidth: bandwidth.Config{
   130  			EgressBitsPerSec:  10,
   131  			IngressBitsPerSec: 50,
   132  			TokenSize:         1,
   133  			Enable:            true,
   134  		},
   135  		Backend: map[string]interface{}{
   136  			"testfs": testfs.Config{Addr: "test-addr", NamePath: namepath.Identity},
   137  		},
   138  	}}, AuthConfig{})
   139  	require.NoError(err)
   140  
   141  	checkBandwidth := func(egress, ingress int64) {
   142  		c, err := m.GetClient("foo")
   143  		require.NoError(err)
   144  		tc, ok := c.(*ThrottledClient)
   145  		require.True(ok)
   146  		require.Equal(egress, tc.EgressLimit())
   147  		require.Equal(ingress, tc.IngressLimit())
   148  	}
   149  
   150  	checkBandwidth(10, 50)
   151  
   152  	watcher := NewBandwidthWatcher(m)
   153  	watcher.Notify(stringset.New("a", "b"))
   154  
   155  	checkBandwidth(5, 25)
   156  }