github.com/wallyworld/juju@v0.0.0-20161013125918-6cf1bc9d917a/worker/uniter/runner/jujuc/testing/storage.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package testing
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"github.com/juju/errors"
    10  	"github.com/juju/testing"
    11  	"github.com/juju/utils/set"
    12  	"gopkg.in/juju/names.v2"
    13  
    14  	"github.com/juju/juju/apiserver/params"
    15  	"github.com/juju/juju/storage"
    16  	"github.com/juju/juju/worker/uniter/runner/jujuc"
    17  )
    18  
    19  // Storage holds the values for the hook context.
    20  type Storage struct {
    21  	Storage    map[names.StorageTag]jujuc.ContextStorageAttachment
    22  	StorageTag names.StorageTag
    23  	Added      map[string]params.StorageConstraints
    24  }
    25  
    26  // SetAttachment adds the attachment to the storage.
    27  func (s *Storage) SetAttachment(attach jujuc.ContextStorageAttachment) {
    28  	if attach == nil || attach == jujuc.ContextStorageAttachment(nil) {
    29  		return
    30  	}
    31  	if s.Storage == nil {
    32  		s.Storage = make(map[names.StorageTag]jujuc.ContextStorageAttachment)
    33  	}
    34  	s.Storage[attach.Tag()] = attach
    35  }
    36  
    37  // SetNewAttachment adds the attachment to the storage.
    38  func (s *Storage) SetNewAttachment(name, location string, kind storage.StorageKind, stub *testing.Stub) {
    39  	tag := names.NewStorageTag(name)
    40  	attachment := &ContextStorageAttachment{
    41  		info: &StorageAttachment{tag, kind, location},
    42  	}
    43  	attachment.stub = stub
    44  	s.SetAttachment(attachment)
    45  }
    46  
    47  // SetBlockStorage adds the attachment to the storage.
    48  func (s *Storage) SetBlockStorage(name, location string, stub *testing.Stub) {
    49  	s.SetNewAttachment(name, location, storage.StorageKindBlock, stub)
    50  }
    51  
    52  // SetStorageTag sets the storage tag to the given ID.
    53  func (s *Storage) SetStorageTag(id string) {
    54  	tag := names.NewStorageTag(id)
    55  	if _, ok := s.Storage[tag]; !ok {
    56  		panic(fmt.Sprintf("storage %q not added yet", id))
    57  	}
    58  	s.StorageTag = tag
    59  }
    60  
    61  // SetUnitStorage sets storage that should be added.
    62  func (s *Storage) SetUnitStorage(name string, constraints params.StorageConstraints) {
    63  	if s.Added == nil {
    64  		s.Added = make(map[string]params.StorageConstraints)
    65  	}
    66  	s.Added[name] = constraints
    67  }
    68  
    69  // AddUnitStorage sets storage that should be added.
    70  func (s *Storage) AddUnitStorage(all map[string]params.StorageConstraints) {
    71  	if s.Added == nil {
    72  		s.Added = make(map[string]params.StorageConstraints)
    73  	}
    74  	for k, v := range all {
    75  		s.Added[k] = v
    76  	}
    77  }
    78  
    79  // ContextStorage is a test double for jujuc.ContextStorage.
    80  type ContextStorage struct {
    81  	contextBase
    82  	info *Storage
    83  }
    84  
    85  // StorageTags implements jujuc.ContextStorage.
    86  func (c *ContextStorage) StorageTags() ([]names.StorageTag, error) {
    87  	c.stub.AddCall("StorageTags")
    88  
    89  	tags := set.NewTags()
    90  	for tag := range c.info.Storage {
    91  		tags.Add(tag)
    92  	}
    93  	storageTags := make([]names.StorageTag, tags.Size())
    94  	for i, tag := range tags.SortedValues() {
    95  		storageTags[i] = tag.(names.StorageTag)
    96  	}
    97  	return storageTags, c.stub.NextErr()
    98  }
    99  
   100  // Storage implements jujuc.ContextStorage.
   101  func (c *ContextStorage) Storage(tag names.StorageTag) (jujuc.ContextStorageAttachment, error) {
   102  	c.stub.AddCall("Storage")
   103  
   104  	storage, ok := c.info.Storage[tag]
   105  	var err error
   106  	if !ok {
   107  		err = errors.NotFoundf("storage")
   108  	}
   109  	return storage, err
   110  
   111  }
   112  
   113  // HookStorage implements jujuc.ContextStorage.
   114  func (c *ContextStorage) HookStorage() (jujuc.ContextStorageAttachment, error) {
   115  	c.stub.AddCall("HookStorage")
   116  
   117  	return c.Storage(c.info.StorageTag)
   118  }
   119  
   120  // AddUnitStorage implements jujuc.ContextStorage.
   121  func (c *ContextStorage) AddUnitStorage(all map[string]params.StorageConstraints) error {
   122  	c.stub.AddCall("AddUnitStorage", all)
   123  	c.info.AddUnitStorage(all)
   124  	return c.stub.NextErr()
   125  }