github.com/mhilton/juju-juju@v0.0.0-20150901100907-a94dd2c73455/state/leadership/util_test.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package leadership_test
     5  
     6  import (
     7  	"fmt"
     8  	"time"
     9  
    10  	"github.com/juju/errors"
    11  	jc "github.com/juju/testing/checkers"
    12  	gc "gopkg.in/check.v1"
    13  
    14  	"github.com/juju/juju/state/lease"
    15  )
    16  
    17  // Client implements lease.Client for testing purposes.
    18  type Client struct {
    19  	leases map[string]lease.Info
    20  	expect []call
    21  	failed string
    22  	done   chan struct{}
    23  }
    24  
    25  // NewClient initializes and returns a new client configured to report
    26  // the supplied leases and expect the supplied calls.
    27  func NewClient(leases map[string]lease.Info, expect []call) *Client {
    28  	if leases == nil {
    29  		leases = make(map[string]lease.Info)
    30  	}
    31  	done := make(chan struct{})
    32  	if len(expect) == 0 {
    33  		close(done)
    34  	}
    35  	return &Client{
    36  		leases: leases,
    37  		expect: expect,
    38  		done:   done,
    39  	}
    40  }
    41  
    42  // Wait will return when all expected calls have been made, or fail the test
    43  // if they don't happen within a second. (You control the clock; your tests
    44  // should pass in *way* less than a second of wall-clock time.)
    45  func (client *Client) Wait(c *gc.C) {
    46  	select {
    47  	case <-client.done:
    48  		if client.failed != "" {
    49  			c.Fatalf(client.failed)
    50  		}
    51  	case <-time.After(time.Second):
    52  		c.Fatalf("Client test took way too long")
    53  	}
    54  }
    55  
    56  // Leases is part of the lease.Client interface.
    57  func (client *Client) Leases() map[string]lease.Info {
    58  	result := make(map[string]lease.Info)
    59  	for k, v := range client.leases {
    60  		result[k] = v
    61  	}
    62  	return result
    63  }
    64  
    65  // call implements the bulk of the lease.Client interface.
    66  func (client *Client) call(method string, args []interface{}) error {
    67  	select {
    68  	case <-client.done:
    69  		return errors.Errorf("Client method called after test complete: %s %v", method, args)
    70  	default:
    71  		defer func() {
    72  			if len(client.expect) == 0 || client.failed != "" {
    73  				close(client.done)
    74  			}
    75  		}()
    76  	}
    77  
    78  	expect := client.expect[0]
    79  	client.expect = client.expect[1:]
    80  	if expect.callback != nil {
    81  		expect.callback(client.leases)
    82  	}
    83  
    84  	if method == expect.method {
    85  		if ok, _ := jc.DeepEqual(args, expect.args); ok {
    86  			return expect.err
    87  		}
    88  	}
    89  	client.failed = fmt.Sprintf("unexpected Client call:\n  actual: %s %v\n  expect: %s %v",
    90  		method, args, expect.method, expect.args,
    91  	)
    92  	return errors.New(client.failed)
    93  }
    94  
    95  // ClaimLease is part of the lease.Client interface.
    96  func (client *Client) ClaimLease(name string, request lease.Request) error {
    97  	return client.call("ClaimLease", []interface{}{name, request})
    98  }
    99  
   100  // ExtendLease is part of the lease.Client interface.
   101  func (client *Client) ExtendLease(name string, request lease.Request) error {
   102  	return client.call("ExtendLease", []interface{}{name, request})
   103  }
   104  
   105  // ExpireLease is part of the lease.Client interface.
   106  func (client *Client) ExpireLease(name string) error {
   107  	return client.call("ExpireLease", []interface{}{name})
   108  }
   109  
   110  // Refresh is part of the lease.Client interface.
   111  func (client *Client) Refresh() error {
   112  	return client.call("Refresh", nil)
   113  }
   114  
   115  // call defines a expected method call on a Client; it encodes:
   116  type call struct {
   117  
   118  	// method is the name of the method.
   119  	method string
   120  
   121  	// args is the expected arguments.
   122  	args []interface{}
   123  
   124  	// err is the error to return.
   125  	err error
   126  
   127  	// callback, if non-nil, will be passed the internal leases dict; for
   128  	// modification, if desired. Otherwise you can use it to, e.g., assert
   129  	// clock time.
   130  	callback func(leases map[string]lease.Info)
   131  }