github.com/mhilton/juju-juju@v0.0.0-20150901100907-a94dd2c73455/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/names"
    10  	"github.com/juju/testing"
    11  	"github.com/juju/utils/set"
    12  
    13  	"github.com/juju/juju/apiserver/params"
    14  	"github.com/juju/juju/storage"
    15  	"github.com/juju/juju/worker/uniter/runner/jujuc"
    16  )
    17  
    18  // Storage holds the values for the hook context.
    19  type Storage struct {
    20  	Storage    map[names.StorageTag]jujuc.ContextStorageAttachment
    21  	StorageTag names.StorageTag
    22  	Added      map[string]params.StorageConstraints
    23  }
    24  
    25  // SetAttachment adds the attachment to the storage.
    26  func (s *Storage) SetAttachment(attach jujuc.ContextStorageAttachment) {
    27  	if attach == nil || attach == jujuc.ContextStorageAttachment(nil) {
    28  		return
    29  	}
    30  	if s.Storage == nil {
    31  		s.Storage = make(map[names.StorageTag]jujuc.ContextStorageAttachment)
    32  	}
    33  	s.Storage[attach.Tag()] = attach
    34  }
    35  
    36  // SetNewAttachment adds the attachment to the storage.
    37  func (s *Storage) SetNewAttachment(name, location string, kind storage.StorageKind, stub *testing.Stub) {
    38  	tag := names.NewStorageTag(name)
    39  	attachment := &ContextStorageAttachment{
    40  		info: &StorageAttachment{tag, kind, location},
    41  	}
    42  	attachment.stub = stub
    43  	s.SetAttachment(attachment)
    44  }
    45  
    46  // SetBlockStorage adds the attachment to the storage.
    47  func (s *Storage) SetBlockStorage(name, location string, stub *testing.Stub) {
    48  	s.SetNewAttachment(name, location, storage.StorageKindBlock, stub)
    49  }
    50  
    51  // SetStorageTag sets the storage tag to the given ID.
    52  func (s *Storage) SetStorageTag(id string) {
    53  	tag := names.NewStorageTag(id)
    54  	if _, ok := s.Storage[tag]; !ok {
    55  		panic(fmt.Sprintf("storage %q not added yet", id))
    56  	}
    57  	s.StorageTag = tag
    58  }
    59  
    60  // SetUnitStorage sets storage that should be added.
    61  func (s *Storage) SetUnitStorage(name string, constraints params.StorageConstraints) {
    62  	if s.Added == nil {
    63  		s.Added = make(map[string]params.StorageConstraints)
    64  	}
    65  	s.Added[name] = constraints
    66  }
    67  
    68  // AddUnitStorage sets storage that should be added.
    69  func (s *Storage) AddUnitStorage(all map[string]params.StorageConstraints) {
    70  	if s.Added == nil {
    71  		s.Added = make(map[string]params.StorageConstraints)
    72  	}
    73  	for k, v := range all {
    74  		s.Added[k] = v
    75  	}
    76  }
    77  
    78  // ContextStorage is a test double for jujuc.ContextStorage.
    79  type ContextStorage struct {
    80  	contextBase
    81  	info *Storage
    82  }
    83  
    84  // StorageTags implements jujuc.ContextStorage.
    85  func (c *ContextStorage) StorageTags() []names.StorageTag {
    86  	c.stub.AddCall("StorageTags")
    87  	c.stub.NextErr()
    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
    98  }
    99  
   100  // Storage implements jujuc.ContextStorage.
   101  func (c *ContextStorage) Storage(tag names.StorageTag) (jujuc.ContextStorageAttachment, bool) {
   102  	c.stub.AddCall("Storage")
   103  	c.stub.NextErr()
   104  
   105  	storage, ok := c.info.Storage[tag]
   106  	return storage, ok
   107  }
   108  
   109  // HookStorage implements jujuc.ContextStorage.
   110  func (c *ContextStorage) HookStorage() (jujuc.ContextStorageAttachment, bool) {
   111  	c.stub.AddCall("HookStorage")
   112  	c.stub.NextErr()
   113  
   114  	return c.Storage(c.info.StorageTag)
   115  }
   116  
   117  // AddUnitStorage implements jujuc.ContextStorage.
   118  func (c *ContextStorage) AddUnitStorage(all map[string]params.StorageConstraints) {
   119  	c.stub.AddCall("AddUnitStorage", all)
   120  	c.stub.NextErr()
   121  
   122  	c.info.AddUnitStorage(all)
   123  }