github.com/wallyworld/juju@v0.0.0-20161013125918-6cf1bc9d917a/state/lease/util_test.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package lease_test
     5  
     6  import (
     7  	"time"
     8  
     9  	jujutxn "github.com/juju/txn"
    10  	"github.com/juju/utils/clock"
    11  	"gopkg.in/mgo.v2"
    12  
    13  	"github.com/juju/juju/mongo"
    14  )
    15  
    16  // Clock exposes time via Now, and can be controlled via Reset and Advance. It
    17  // can be configured to Advance automatically whenever Now is called. Attempts
    18  // to call Alarm will panic: they're not useful to a clock.Client itself, but
    19  // are extremely helpful when driving one.
    20  type Clock struct {
    21  	clock.Clock
    22  	now  time.Time
    23  	step time.Duration
    24  }
    25  
    26  // NewClock returns a *Clock, set to now, that advances by step whenever Now()
    27  // is called.
    28  func NewClock(now time.Time, step time.Duration) *Clock {
    29  	return &Clock{now: now, step: step}
    30  }
    31  
    32  // Now is part of the clock.Clock interface.
    33  func (clock *Clock) Now() time.Time {
    34  	defer clock.Advance(clock.step)
    35  	return clock.now
    36  }
    37  
    38  // Reset causes the clock to act as though it had just been created with
    39  // NewClock(now, step).
    40  func (clock *Clock) Reset(now time.Time, step time.Duration) {
    41  	clock.now = now
    42  	clock.step = step
    43  }
    44  
    45  // Advance advances the clock by the supplied time.
    46  func (clock *Clock) Advance(duration time.Duration) {
    47  	clock.now = clock.now.Add(duration)
    48  }
    49  
    50  // Mongo exposes database operations. It uses a real database -- we can't mock
    51  // mongo out, we need to check it really actually works -- but it's good to
    52  // have the runner accessible for adversarial transaction tests.
    53  type Mongo struct {
    54  	database *mgo.Database
    55  	runner   jujutxn.Runner
    56  }
    57  
    58  // NewMongo returns a *Mongo backed by the supplied database.
    59  func NewMongo(database *mgo.Database) *Mongo {
    60  	return &Mongo{
    61  		database: database,
    62  		runner: jujutxn.NewRunner(jujutxn.RunnerParams{
    63  			Database: database,
    64  		}),
    65  	}
    66  }
    67  
    68  // GetCollection is part of the lease.Mongo interface.
    69  func (m *Mongo) GetCollection(name string) (mongo.Collection, func()) {
    70  	return mongo.CollectionFromName(m.database, name)
    71  }
    72  
    73  // RunTransaction is part of the lease.Mongo interface.
    74  func (m *Mongo) RunTransaction(getTxn jujutxn.TransactionSource) error {
    75  	return m.runner.Run(getTxn)
    76  }