github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/worker/uniter/runner/jujuc/jujuctesting/storage.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package jujuctesting
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"github.com/juju/errors"
    10  	"github.com/juju/testing"
    11  	"gopkg.in/juju/names.v2"
    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, error) {
    86  	c.stub.AddCall("StorageTags")
    87  
    88  	tags := names.NewSet()
    89  	for tag := range c.info.Storage {
    90  		tags.Add(tag)
    91  	}
    92  	storageTags := make([]names.StorageTag, tags.Size())
    93  	for i, tag := range tags.SortedValues() {
    94  		storageTags[i] = tag.(names.StorageTag)
    95  	}
    96  	return storageTags, c.stub.NextErr()
    97  }
    98  
    99  // Storage implements jujuc.ContextStorage.
   100  func (c *ContextStorage) Storage(tag names.StorageTag) (jujuc.ContextStorageAttachment, error) {
   101  	c.stub.AddCall("Storage")
   102  
   103  	storage, ok := c.info.Storage[tag]
   104  	var err error
   105  	if !ok {
   106  		err = errors.NotFoundf("storage")
   107  	}
   108  	return storage, err
   109  
   110  }
   111  
   112  // HookStorage implements jujuc.ContextStorage.
   113  func (c *ContextStorage) HookStorage() (jujuc.ContextStorageAttachment, error) {
   114  	c.stub.AddCall("HookStorage")
   115  
   116  	return c.Storage(c.info.StorageTag)
   117  }
   118  
   119  // AddUnitStorage implements jujuc.ContextStorage.
   120  func (c *ContextStorage) AddUnitStorage(all map[string]params.StorageConstraints) error {
   121  	c.stub.AddCall("AddUnitStorage", all)
   122  	c.info.AddUnitStorage(all)
   123  	return c.stub.NextErr()
   124  }