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