github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/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  type ContextSuite struct {
    34  	jujuctesting.ContextSuite
    35  	testing.BaseSuite
    36  }
    37  
    38  func (s *ContextSuite) SetUpTest(c *gc.C) {
    39  	s.ContextSuite.SetUpTest(c)
    40  	s.BaseSuite.SetUpTest(c)
    41  }
    42  
    43  func (s *ContextSuite) newHookContext(c *gc.C) *Context {
    44  	hctx, info := s.ContextSuite.NewHookContext()
    45  	return &Context{
    46  		Context: hctx,
    47  		info:    info,
    48  	}
    49  }
    50  
    51  func (s *ContextSuite) GetHookContext(c *gc.C, relid int, remote string) *Context {
    52  	c.Assert(relid, gc.Equals, -1)
    53  	return s.newHookContext(c)
    54  }
    55  
    56  func (s *ContextSuite) GetStatusHookContext(c *gc.C) *Context {
    57  	return s.newHookContext(c)
    58  }
    59  
    60  type Context struct {
    61  	jujuc.Context
    62  	info *jujuctesting.ContextInfo
    63  
    64  	metrics        []jujuc.Metric
    65  	canAddMetrics  bool
    66  	rebootPriority jujuc.RebootPriority
    67  	shouldError    bool
    68  }
    69  
    70  func (c *Context) AddMetric(key, value string, created time.Time) error {
    71  	if !c.canAddMetrics {
    72  		return fmt.Errorf("metrics disabled")
    73  	}
    74  	c.metrics = append(c.metrics, jujuc.Metric{
    75  		Key:   key,
    76  		Value: value,
    77  		Time:  created,
    78  	})
    79  	return c.Context.AddMetric(key, value, created)
    80  }
    81  
    82  func (c *Context) AddMetricLabels(key, value string, created time.Time, labels map[string]string) error {
    83  	if !c.canAddMetrics {
    84  		return fmt.Errorf("metrics disabled")
    85  	}
    86  	c.metrics = append(c.metrics, jujuc.Metric{
    87  		Key:    key,
    88  		Value:  value,
    89  		Time:   created,
    90  		Labels: labels,
    91  	})
    92  	return c.Context.AddMetricLabels(key, value, created, labels)
    93  }
    94  
    95  func (c *Context) RequestReboot(priority jujuc.RebootPriority) error {
    96  	c.rebootPriority = priority
    97  	if c.shouldError {
    98  		return fmt.Errorf("RequestReboot error!")
    99  	} else {
   100  		return nil
   101  	}
   102  }