github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/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  	"github.com/juju/clock"
    10  	jujutxn "github.com/juju/txn"
    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  // This differs from github.com/juju/testclock.Clock in that it has a Reset() function.
    21  type Clock struct {
    22  	clock.Clock
    23  	now time.Time
    24  }
    25  
    26  // NewClock returns a *Clock, set to now.
    27  func NewClock(now time.Time) *Clock {
    28  	return &Clock{now: now}
    29  }
    30  
    31  // Now is part of the clock.Clock interface.
    32  func (clock *Clock) Now() time.Time {
    33  	return clock.now
    34  }
    35  
    36  // Reset causes the clock to act as though it had just been created with
    37  // NewClock(now).
    38  func (clock *Clock) Reset(now time.Time) {
    39  	clock.now = now
    40  }
    41  
    42  // Advance advances the clock by the supplied time.
    43  func (clock *Clock) Advance(duration time.Duration) {
    44  	clock.now = clock.now.Add(duration)
    45  }
    46  
    47  type GlobalClock struct {
    48  	*Clock
    49  }
    50  
    51  func (clock GlobalClock) Now() (time.Time, error) {
    52  	return clock.Clock.Now(), nil
    53  }
    54  
    55  // Mongo exposes database operations. It uses a real database -- we can't mock
    56  // mongo out, we need to check it really actually works -- but it's good to
    57  // have the runner accessible for adversarial transaction tests.
    58  type Mongo struct {
    59  	database *mgo.Database
    60  	runner   jujutxn.Runner
    61  }
    62  
    63  // NewMongo returns a *Mongo backed by the supplied database.
    64  func NewMongo(database *mgo.Database) *Mongo {
    65  	return &Mongo{
    66  		database: database,
    67  		runner: jujutxn.NewRunner(jujutxn.RunnerParams{
    68  			Database: database,
    69  		}),
    70  	}
    71  }
    72  
    73  // GetCollection is part of the lease.Mongo interface.
    74  func (m *Mongo) GetCollection(name string) (mongo.Collection, func()) {
    75  	return mongo.CollectionFromName(m.database, name)
    76  }
    77  
    78  // RunTransaction is part of the lease.Mongo interface.
    79  func (m *Mongo) RunTransaction(getTxn jujutxn.TransactionSource) error {
    80  	return m.runner.Run(getTxn)
    81  }