github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/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  	"github.com/juju/juju/worker/uniter/runner/jujuc/jujuctesting"
    18  )
    19  
    20  const (
    21  	formatYaml = iota
    22  	formatJson
    23  )
    24  
    25  func bufferBytes(stream io.Writer) []byte {
    26  	return stream.(*bytes.Buffer).Bytes()
    27  }
    28  
    29  func bufferString(w io.Writer) string {
    30  	return w.(*bytes.Buffer).String()
    31  }
    32  
    33  func cmdString(cmd string) string {
    34  	return cmd + jujuc.CmdSuffix
    35  }
    36  
    37  type ContextSuite struct {
    38  	jujuctesting.ContextSuite
    39  	testing.BaseSuite
    40  
    41  	rels map[int]*jujuctesting.ContextRelation
    42  }
    43  
    44  func (s *ContextSuite) SetUpTest(c *gc.C) {
    45  	s.ContextSuite.SetUpTest(c)
    46  	s.BaseSuite.SetUpTest(c)
    47  }
    48  
    49  func (s *ContextSuite) newHookContext(c *gc.C) *Context {
    50  	hctx, info := s.ContextSuite.NewHookContext()
    51  	return &Context{
    52  		Context: hctx,
    53  		info:    info,
    54  	}
    55  }
    56  
    57  func (s *ContextSuite) GetHookContext(c *gc.C, relid int, remote string) *Context {
    58  	c.Assert(relid, gc.Equals, -1)
    59  	return s.newHookContext(c)
    60  }
    61  
    62  func (s *ContextSuite) GetStatusHookContext(c *gc.C) *Context {
    63  	return s.newHookContext(c)
    64  }
    65  
    66  type Context struct {
    67  	jujuc.Context
    68  	info *jujuctesting.ContextInfo
    69  
    70  	metrics        []jujuc.Metric
    71  	canAddMetrics  bool
    72  	rebootPriority jujuc.RebootPriority
    73  	shouldError    bool
    74  }
    75  
    76  func (c *Context) AddMetric(key, value string, created time.Time) error {
    77  	if !c.canAddMetrics {
    78  		return fmt.Errorf("metrics disabled")
    79  	}
    80  	c.metrics = append(c.metrics, jujuc.Metric{
    81  		Key:   key,
    82  		Value: value,
    83  		Time:  created,
    84  	})
    85  	return c.Context.AddMetric(key, value, created)
    86  }
    87  
    88  func (c *Context) AddMetricLabels(key, value string, created time.Time, labels map[string]string) error {
    89  	if !c.canAddMetrics {
    90  		return fmt.Errorf("metrics disabled")
    91  	}
    92  	c.metrics = append(c.metrics, jujuc.Metric{
    93  		Key:    key,
    94  		Value:  value,
    95  		Time:   created,
    96  		Labels: labels,
    97  	})
    98  	return c.Context.AddMetricLabels(key, value, created, labels)
    99  }
   100  
   101  func (c *Context) RequestReboot(priority jujuc.RebootPriority) error {
   102  	c.rebootPriority = priority
   103  	if c.shouldError {
   104  		return fmt.Errorf("RequestReboot error!")
   105  	} else {
   106  		return nil
   107  	}
   108  }