github.com/makyo/juju@v0.0.0-20160425123129-2608902037e9/state/statetest/persistence_stubs.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package statetest
     5  
     6  import (
     7  	"reflect"
     8  
     9  	"github.com/juju/errors"
    10  	"github.com/juju/testing"
    11  	jujutxn "github.com/juju/txn"
    12  	"gopkg.in/mgo.v2/txn"
    13  )
    14  
    15  type StubPersistence struct {
    16  	*testing.Stub
    17  
    18  	RunFunc func(jujutxn.TransactionSource) error
    19  
    20  	ReturnAll interface{} // homegenous(?) list of doc struct (not pointers)
    21  	ReturnOne interface{} // a doc struct (not a pointer)
    22  
    23  	ReturnServiceExistsOps           []txn.Op
    24  	ReturnIncCharmModifiedVersionOps []txn.Op
    25  	ReturnNewCleanupOp               *txn.Op
    26  }
    27  
    28  func NewStubPersistence(stub *testing.Stub) *StubPersistence {
    29  	s := &StubPersistence{
    30  		Stub: stub,
    31  	}
    32  	s.RunFunc = s.run
    33  	return s
    34  }
    35  
    36  func (s *StubPersistence) One(collName, id string, doc interface{}) error {
    37  	s.AddCall("One", collName, id, doc)
    38  	if err := s.NextErr(); err != nil {
    39  		return errors.Trace(err)
    40  	}
    41  
    42  	if reflect.TypeOf(s.ReturnOne) == nil {
    43  		return errors.NotFoundf("resource")
    44  	}
    45  	ptr := reflect.ValueOf(doc)
    46  	newVal := reflect.ValueOf(s.ReturnOne)
    47  	ptr.Elem().Set(newVal)
    48  	return nil
    49  }
    50  
    51  func (s *StubPersistence) All(collName string, query, docs interface{}) error {
    52  	s.AddCall("All", collName, query, docs)
    53  	if err := s.NextErr(); err != nil {
    54  		return errors.Trace(err)
    55  	}
    56  
    57  	ptr := reflect.ValueOf(docs)
    58  	if reflect.TypeOf(s.ReturnAll) == nil {
    59  		ptr.Elem().SetLen(0)
    60  	} else {
    61  		newVal := reflect.ValueOf(s.ReturnAll)
    62  		ptr.Elem().Set(newVal)
    63  	}
    64  	return nil
    65  }
    66  
    67  func (s *StubPersistence) Run(buildTxn jujutxn.TransactionSource) error {
    68  	s.AddCall("Run", buildTxn)
    69  	if err := s.NextErr(); err != nil {
    70  		return errors.Trace(err)
    71  	}
    72  
    73  	if err := s.run(buildTxn); err != nil {
    74  		return errors.Trace(err)
    75  	}
    76  
    77  	return nil
    78  }
    79  
    80  // See github.com/juju/txn.transactionRunner.Run.
    81  func (s *StubPersistence) run(buildTxn jujutxn.TransactionSource) error {
    82  	for i := 0; ; i++ {
    83  		ops, err := buildTxn(i)
    84  		if errors.Cause(err) == jujutxn.ErrTransientFailure {
    85  			continue
    86  		}
    87  		if errors.Cause(err) == jujutxn.ErrNoOperations {
    88  			return nil
    89  		}
    90  		if err != nil {
    91  			return err
    92  		}
    93  
    94  		err = s.RunTransaction(ops)
    95  		if errors.Cause(err) == txn.ErrAborted {
    96  			continue
    97  		}
    98  		if err != nil {
    99  			return err
   100  		}
   101  		return nil
   102  	}
   103  	return nil
   104  }
   105  
   106  func (s *StubPersistence) RunTransaction(ops []txn.Op) error {
   107  	s.AddCall("RunTransaction", ops)
   108  	if err := s.NextErr(); err != nil {
   109  		return errors.Trace(err)
   110  	}
   111  
   112  	return nil
   113  }
   114  
   115  func (s *StubPersistence) ServiceExistsOps(serviceID string) []txn.Op {
   116  	s.AddCall("ServiceExistsOps", serviceID)
   117  	// pop off an error so num errors == num calls, even though this call
   118  	// doesn't actually use the error.
   119  	s.NextErr()
   120  
   121  	return s.ReturnServiceExistsOps
   122  }
   123  
   124  func (s *StubPersistence) IncCharmModifiedVersionOps(serviceID string) []txn.Op {
   125  	s.AddCall("IncCharmModifiedVersionOps", serviceID)
   126  	// pop off an error so num errors == num calls, even though this call
   127  	// doesn't actually use the error.
   128  	s.NextErr()
   129  
   130  	return s.ReturnIncCharmModifiedVersionOps
   131  }
   132  
   133  func (s *StubPersistence) NewCleanupOp(kind, prefix string) txn.Op {
   134  	s.AddCall("NewCleanupOp", kind, prefix)
   135  	// pop off an error so num errors == num calls, even though this call
   136  	// doesn't actually use the error.
   137  	s.NextErr()
   138  
   139  	return *s.ReturnNewCleanupOp
   140  }