github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/state/bakerystorage/storage_test.go (about)

     1  // Copyright 2014-2016 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package bakerystorage
     5  
     6  import (
     7  	"encoding/json"
     8  	"time" // Only used for time types.
     9  
    10  	gitjujutesting "github.com/juju/testing"
    11  	jc "github.com/juju/testing/checkers"
    12  	gc "gopkg.in/check.v1"
    13  	"gopkg.in/macaroon-bakery.v2-unstable/bakery"
    14  	"gopkg.in/macaroon-bakery.v2-unstable/bakery/mgostorage"
    15  	"gopkg.in/mgo.v2"
    16  
    17  	"github.com/juju/juju/mongo"
    18  	"github.com/juju/juju/testing"
    19  )
    20  
    21  type StorageSuite struct {
    22  	testing.BaseSuite
    23  	gitjujutesting.Stub
    24  	collection      mockCollection
    25  	memStorage      bakery.Storage
    26  	closeCollection func()
    27  	config          Config
    28  }
    29  
    30  var _ = gc.Suite(&StorageSuite{})
    31  
    32  func (s *StorageSuite) SetUpTest(c *gc.C) {
    33  	s.BaseSuite.SetUpTest(c)
    34  	s.Stub.ResetCalls()
    35  	s.collection = mockCollection{Stub: &s.Stub}
    36  	s.closeCollection = func() {
    37  		s.AddCall("Close")
    38  		s.PopNoErr()
    39  	}
    40  	s.memStorage = bakery.NewMemStorage()
    41  	s.config = Config{
    42  		GetCollection: func() (mongo.Collection, func()) {
    43  			s.AddCall("GetCollection")
    44  			s.PopNoErr()
    45  			return &s.collection, s.closeCollection
    46  		},
    47  		GetStorage: func(rootKeys *mgostorage.RootKeys, coll mongo.Collection, expireAfter time.Duration) bakery.Storage {
    48  			s.AddCall("GetStorage", coll, expireAfter)
    49  			s.PopNoErr()
    50  			return s.memStorage
    51  		},
    52  	}
    53  }
    54  
    55  func (s *StorageSuite) TestValidateConfigGetCollection(c *gc.C) {
    56  	s.config.GetCollection = nil
    57  	_, err := New(s.config)
    58  	c.Assert(err, gc.ErrorMatches, "validating config: nil GetCollection not valid")
    59  }
    60  
    61  func (s *StorageSuite) TestValidateConfigGetStorage(c *gc.C) {
    62  	s.config.GetStorage = nil
    63  	_, err := New(s.config)
    64  	c.Assert(err, gc.ErrorMatches, "validating config: nil GetStorage not valid")
    65  }
    66  
    67  func (s *StorageSuite) TestExpireAfter(c *gc.C) {
    68  	store, err := New(s.config)
    69  	c.Assert(err, jc.ErrorIsNil)
    70  
    71  	store = store.ExpireAfter(24 * time.Hour)
    72  	c.Assert(ExpireAfter(store), gc.Equals, 24*time.Hour)
    73  }
    74  
    75  func (s *StorageSuite) TestGet(c *gc.C) {
    76  	store, err := New(s.config)
    77  	c.Assert(err, jc.ErrorIsNil)
    78  
    79  	rootKey, id, err := store.RootKey()
    80  	c.Assert(err, jc.ErrorIsNil)
    81  
    82  	item, err := store.Get(id)
    83  	c.Assert(err, jc.ErrorIsNil)
    84  	c.Assert(item, jc.DeepEquals, rootKey)
    85  	s.CheckCalls(c, []gitjujutesting.StubCall{
    86  		{"GetCollection", nil},
    87  		{"GetStorage", []interface{}{&s.collection, time.Duration(0)}},
    88  		{"Close", nil},
    89  		{"GetCollection", nil},
    90  		{"GetStorage", []interface{}{&s.collection, time.Duration(0)}},
    91  		{"Close", nil},
    92  	})
    93  }
    94  
    95  func (s *StorageSuite) TestGetNotFound(c *gc.C) {
    96  	store, err := New(s.config)
    97  	c.Assert(err, jc.ErrorIsNil)
    98  	c.Log("1.")
    99  	_, err = store.Get([]byte("foo"))
   100  	c.Log("2.")
   101  	c.Assert(err, gc.Equals, bakery.ErrNotFound)
   102  }
   103  
   104  func (s *StorageSuite) TestGetLegacyFallback(c *gc.C) {
   105  	store, err := New(s.config)
   106  	c.Assert(err, jc.ErrorIsNil)
   107  
   108  	var rk legacyRootKey
   109  	err = json.Unmarshal([]byte("{\"RootKey\":\"ibbhlQv5+yf7UMNI77W4hxQeQjRdMxs0\"}"), &rk)
   110  	c.Assert(err, jc.ErrorIsNil)
   111  
   112  	item, err := store.Get([]byte("oldkey"))
   113  	c.Assert(err, jc.ErrorIsNil)
   114  	c.Assert(item, jc.DeepEquals, rk.RootKey)
   115  	s.CheckCalls(c, []gitjujutesting.StubCall{
   116  		{"GetCollection", nil},
   117  		{"GetStorage", []interface{}{&s.collection, time.Duration(0)}},
   118  		{"GetCollection", nil},
   119  		{"FindId", []interface{}{"oldkey"}},
   120  		{"One", []interface{}{&storageDoc{
   121  			// Set by mock, not in input. Unimportant anyway.
   122  			Location: "oldkey",
   123  			Item:     "{\"RootKey\":\"ibbhlQv5+yf7UMNI77W4hxQeQjRdMxs0\"}",
   124  		}}},
   125  		{"Close", nil},
   126  		{"Close", nil},
   127  	})
   128  }
   129  
   130  type mockCollection struct {
   131  	mongo.WriteCollection
   132  	*gitjujutesting.Stub
   133  }
   134  
   135  func (c *mockCollection) FindId(id interface{}) mongo.Query {
   136  	c.MethodCall(c, "FindId", id)
   137  	c.PopNoErr()
   138  	return &mockQuery{Stub: c.Stub, id: id}
   139  }
   140  
   141  func (c *mockCollection) Writeable() mongo.WriteCollection {
   142  	c.MethodCall(c, "Writeable")
   143  	c.PopNoErr()
   144  	return c
   145  }
   146  
   147  type mockQuery struct {
   148  	mongo.Query
   149  	*gitjujutesting.Stub
   150  	id interface{}
   151  }
   152  
   153  func (q *mockQuery) One(result interface{}) error {
   154  	q.MethodCall(q, "One", result)
   155  
   156  	location := q.id.(string)
   157  	if location != "oldkey" {
   158  		return mgo.ErrNotFound
   159  	}
   160  	*result.(*storageDoc) = storageDoc{
   161  		Location: q.id.(string),
   162  		Item:     "{\"RootKey\":\"ibbhlQv5+yf7UMNI77W4hxQeQjRdMxs0\"}",
   163  	}
   164  	return q.NextErr()
   165  }