github.com/juju/juju@v0.0.0-20240327075706-a90865de2538/worker/retrystrategy/fixture_test.go (about)

     1  // Copyright 2016 Canonical Ltd.
     2  // Copyright 2016 Cloudbase Solutions
     3  // Licensed under the AGPLv3, see LICENCE file for details.
     4  
     5  package retrystrategy_test
     6  
     7  import (
     8  	"time"
     9  
    10  	"github.com/juju/loggo"
    11  	"github.com/juju/names/v5"
    12  	"github.com/juju/testing"
    13  	jc "github.com/juju/testing/checkers"
    14  	"github.com/juju/worker/v3"
    15  	"github.com/juju/worker/v3/workertest"
    16  	gc "gopkg.in/check.v1"
    17  
    18  	"github.com/juju/juju/core/watcher"
    19  	"github.com/juju/juju/rpc/params"
    20  	coretesting "github.com/juju/juju/testing"
    21  	"github.com/juju/juju/worker/retrystrategy"
    22  )
    23  
    24  type fixture struct {
    25  	testing.Stub
    26  }
    27  
    28  func newFixture(c *gc.C, errs ...error) *fixture {
    29  	fix := &fixture{}
    30  	c.Assert(nil, jc.ErrorIsNil)
    31  	fix.SetErrors(errs...)
    32  	return fix
    33  }
    34  
    35  func (fix *fixture) Run(c *gc.C, test func(worker.Worker)) {
    36  	stubRetryStrategy := params.RetryStrategy{
    37  		ShouldRetry: true,
    38  	}
    39  	stubTag := stubTag{}
    40  	stubFacade := newStubFacade(c, &fix.Stub, stubRetryStrategy, stubTag)
    41  	stubConfig := retrystrategy.WorkerConfig{
    42  		Facade:        stubFacade,
    43  		AgentTag:      stubTag,
    44  		RetryStrategy: stubRetryStrategy,
    45  		Logger:        loggo.GetLogger("test"),
    46  	}
    47  
    48  	w, err := retrystrategy.NewRetryStrategyWorker(stubConfig)
    49  	c.Assert(err, jc.ErrorIsNil)
    50  	done := make(chan struct{})
    51  	go func() {
    52  		defer close(done)
    53  		defer worker.Stop(w)
    54  		test(w)
    55  	}()
    56  
    57  	select {
    58  	case <-done:
    59  	case <-time.After(coretesting.LongWait):
    60  		c.Fatalf("test timed out")
    61  	}
    62  }
    63  
    64  type stubFacade struct {
    65  	c               *gc.C
    66  	stub            *testing.Stub
    67  	watcher         *stubWatcher
    68  	count           int
    69  	initialStrategy params.RetryStrategy
    70  	stubTag         names.Tag
    71  }
    72  
    73  func newStubFacade(c *gc.C, stub *testing.Stub, initialStrategy params.RetryStrategy, stubTag names.Tag) *stubFacade {
    74  	return &stubFacade{
    75  		c:               c,
    76  		stub:            stub,
    77  		watcher:         newStubWatcher(),
    78  		count:           0,
    79  		initialStrategy: initialStrategy,
    80  		stubTag:         stubTag,
    81  	}
    82  }
    83  
    84  // WatchRetryStrategy is part of the retrystrategy Facade
    85  func (f *stubFacade) WatchRetryStrategy(agentTag names.Tag) (watcher.NotifyWatcher, error) {
    86  	f.c.Assert(agentTag, gc.Equals, f.stubTag)
    87  	f.stub.AddCall("WatchRetryStrategy", agentTag)
    88  	err := f.stub.NextErr()
    89  	if err != nil {
    90  		return nil, err
    91  	}
    92  	return f.watcher, nil
    93  }
    94  
    95  // RetryStrategy is part of the retrystrategy Facade
    96  func (f *stubFacade) RetryStrategy(agentTag names.Tag) (params.RetryStrategy, error) {
    97  	f.c.Assert(agentTag, gc.Equals, f.stubTag)
    98  	f.stub.AddCall("RetryStrategy", agentTag)
    99  	f.count = f.count + 1
   100  	// Change the strategy after 2 handles
   101  	if f.count == 2 {
   102  		f.initialStrategy.ShouldRetry = !f.initialStrategy.ShouldRetry
   103  	}
   104  	return f.initialStrategy, f.stub.NextErr()
   105  }
   106  
   107  type stubWatcher struct {
   108  	worker.Worker
   109  	notifyChan <-chan struct{}
   110  }
   111  
   112  func newStubWatcher() *stubWatcher {
   113  	changes := make(chan struct{}, 3)
   114  	changes <- struct{}{}
   115  	changes <- struct{}{}
   116  	changes <- struct{}{}
   117  	return &stubWatcher{
   118  		Worker:     workertest.NewErrorWorker(nil),
   119  		notifyChan: changes,
   120  	}
   121  }
   122  
   123  // Changes is part of the watcher.NotifyWatcher interface
   124  func (w *stubWatcher) Changes() watcher.NotifyChannel {
   125  	return w.notifyChan
   126  }
   127  
   128  type stubTag struct {
   129  	names.Tag
   130  }