github.com/wallyworld/juju@v0.0.0-20161013125918-6cf1bc9d917a/worker/uniter/runner/jujuc/util_test.go (about)

     1  // Copyright 2012, 2013 Canonical Ltd.
     2  // Copyright 2014 Cloudbase Solutions SRL
     3  // Licensed under the AGPLv3, see LICENCE file for details.
     4  
     5  package jujuc_test
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"time"
    12  
    13  	gc "gopkg.in/check.v1"
    14  
    15  	"github.com/juju/juju/testing"
    16  	"github.com/juju/juju/worker/uniter/runner/jujuc"
    17  	jujuctesting "github.com/juju/juju/worker/uniter/runner/jujuc/testing"
    18  )
    19  
    20  func bufferBytes(stream io.Writer) []byte {
    21  	return stream.(*bytes.Buffer).Bytes()
    22  }
    23  
    24  func bufferString(w io.Writer) string {
    25  	return w.(*bytes.Buffer).String()
    26  }
    27  
    28  func cmdString(cmd string) string {
    29  	return cmd + jujuc.CmdSuffix
    30  }
    31  
    32  type ContextSuite struct {
    33  	jujuctesting.ContextSuite
    34  	testing.BaseSuite
    35  
    36  	rels map[int]*jujuctesting.ContextRelation
    37  }
    38  
    39  func (s *ContextSuite) SetUpTest(c *gc.C) {
    40  	s.ContextSuite.SetUpTest(c)
    41  	s.BaseSuite.SetUpTest(c)
    42  }
    43  
    44  func (s *ContextSuite) newHookContext(c *gc.C) *Context {
    45  	hctx, info := s.ContextSuite.NewHookContext()
    46  	return &Context{
    47  		Context: hctx,
    48  		info:    info,
    49  	}
    50  }
    51  
    52  func (s *ContextSuite) GetHookContext(c *gc.C, relid int, remote string) *Context {
    53  	c.Assert(relid, gc.Equals, -1)
    54  	return s.newHookContext(c)
    55  }
    56  
    57  func (s *ContextSuite) GetStatusHookContext(c *gc.C) *Context {
    58  	return s.newHookContext(c)
    59  }
    60  
    61  type Context struct {
    62  	jujuc.Context
    63  	info *jujuctesting.ContextInfo
    64  
    65  	metrics        []jujuc.Metric
    66  	canAddMetrics  bool
    67  	rebootPriority jujuc.RebootPriority
    68  	shouldError    bool
    69  }
    70  
    71  func (c *Context) AddMetric(key, value string, created time.Time) error {
    72  	if !c.canAddMetrics {
    73  		return fmt.Errorf("metrics disabled")
    74  	}
    75  	c.metrics = append(c.metrics, jujuc.Metric{
    76  		Key:   key,
    77  		Value: value,
    78  		Time:  created,
    79  	})
    80  	return c.Context.AddMetric(key, value, created)
    81  }
    82  
    83  func (c *Context) RequestReboot(priority jujuc.RebootPriority) error {
    84  	c.rebootPriority = priority
    85  	if c.shouldError {
    86  		return fmt.Errorf("RequestReboot error!")
    87  	} else {
    88  		return nil
    89  	}
    90  }