github.com/wfusion/gofusion@v1.1.14/common/infra/asynq/pkg/testbroker/testbroker.go (about)

     1  // Copyright 2020 Kentaro Hibino. All rights reserved.
     2  // Use of this source code is governed by a MIT license
     3  // that can be found in the LICENSE file.
     4  
     5  // Package testbroker exports a broker implementation that should be used in package testing.
     6  package testbroker
     7  
     8  import (
     9  	"context"
    10  	"errors"
    11  	"sync"
    12  	"time"
    13  
    14  	"github.com/redis/go-redis/v9"
    15  
    16  	"github.com/wfusion/gofusion/common/infra/asynq/pkg/base"
    17  )
    18  
    19  var errRedisDown = errors.New("testutil: redis is down")
    20  
    21  // TestBroker is a broker implementation which enables
    22  // to simulate Redis failure in tests.
    23  type TestBroker struct {
    24  	mu       sync.Mutex
    25  	sleeping bool
    26  
    27  	// real broker
    28  	real base.Broker
    29  }
    30  
    31  // Make sure TestBroker implements Broker interface at compile time.
    32  var _ base.Broker = (*TestBroker)(nil)
    33  
    34  func NewTestBroker(b base.Broker) *TestBroker {
    35  	return &TestBroker{real: b}
    36  }
    37  
    38  func (tb *TestBroker) Sleep() {
    39  	tb.mu.Lock()
    40  	defer tb.mu.Unlock()
    41  	tb.sleeping = true
    42  }
    43  
    44  func (tb *TestBroker) Wakeup() {
    45  	tb.mu.Lock()
    46  	defer tb.mu.Unlock()
    47  	tb.sleeping = false
    48  }
    49  
    50  func (tb *TestBroker) Enqueue(ctx context.Context, msg *base.TaskMessage) error {
    51  	tb.mu.Lock()
    52  	defer tb.mu.Unlock()
    53  	if tb.sleeping {
    54  		return errRedisDown
    55  	}
    56  	return tb.real.Enqueue(ctx, msg)
    57  }
    58  
    59  func (tb *TestBroker) EnqueueUnique(ctx context.Context, msg *base.TaskMessage, ttl time.Duration) error {
    60  	tb.mu.Lock()
    61  	defer tb.mu.Unlock()
    62  	if tb.sleeping {
    63  		return errRedisDown
    64  	}
    65  	return tb.real.EnqueueUnique(ctx, msg, ttl)
    66  }
    67  
    68  func (tb *TestBroker) Dequeue(qnames ...string) (*base.TaskMessage, time.Time, error) {
    69  	tb.mu.Lock()
    70  	defer tb.mu.Unlock()
    71  	if tb.sleeping {
    72  		return nil, time.Time{}, errRedisDown
    73  	}
    74  	return tb.real.Dequeue(qnames...)
    75  }
    76  
    77  func (tb *TestBroker) Done(ctx context.Context, msg *base.TaskMessage) error {
    78  	tb.mu.Lock()
    79  	defer tb.mu.Unlock()
    80  	if tb.sleeping {
    81  		return errRedisDown
    82  	}
    83  	return tb.real.Done(ctx, msg)
    84  }
    85  
    86  func (tb *TestBroker) MarkAsComplete(ctx context.Context, msg *base.TaskMessage) error {
    87  	tb.mu.Lock()
    88  	defer tb.mu.Unlock()
    89  	if tb.sleeping {
    90  		return errRedisDown
    91  	}
    92  	return tb.real.MarkAsComplete(ctx, msg)
    93  }
    94  
    95  func (tb *TestBroker) Requeue(ctx context.Context, msg *base.TaskMessage) error {
    96  	tb.mu.Lock()
    97  	defer tb.mu.Unlock()
    98  	if tb.sleeping {
    99  		return errRedisDown
   100  	}
   101  	return tb.real.Requeue(ctx, msg)
   102  }
   103  
   104  func (tb *TestBroker) Schedule(ctx context.Context, msg *base.TaskMessage, processAt time.Time) error {
   105  	tb.mu.Lock()
   106  	defer tb.mu.Unlock()
   107  	if tb.sleeping {
   108  		return errRedisDown
   109  	}
   110  	return tb.real.Schedule(ctx, msg, processAt)
   111  }
   112  
   113  func (tb *TestBroker) ScheduleUnique(ctx context.Context, msg *base.TaskMessage,
   114  	processAt time.Time, ttl time.Duration) error {
   115  	tb.mu.Lock()
   116  	defer tb.mu.Unlock()
   117  	if tb.sleeping {
   118  		return errRedisDown
   119  	}
   120  	return tb.real.ScheduleUnique(ctx, msg, processAt, ttl)
   121  }
   122  
   123  func (tb *TestBroker) Retry(ctx context.Context, msg *base.TaskMessage,
   124  	processAt time.Time, errMsg string, isFailure bool) error {
   125  	tb.mu.Lock()
   126  	defer tb.mu.Unlock()
   127  	if tb.sleeping {
   128  		return errRedisDown
   129  	}
   130  	return tb.real.Retry(ctx, msg, processAt, errMsg, isFailure)
   131  }
   132  
   133  func (tb *TestBroker) Archive(ctx context.Context, msg *base.TaskMessage, errMsg string) error {
   134  	tb.mu.Lock()
   135  	defer tb.mu.Unlock()
   136  	if tb.sleeping {
   137  		return errRedisDown
   138  	}
   139  	return tb.real.Archive(ctx, msg, errMsg)
   140  }
   141  
   142  func (tb *TestBroker) ForwardIfReady(qnames ...string) error {
   143  	tb.mu.Lock()
   144  	defer tb.mu.Unlock()
   145  	if tb.sleeping {
   146  		return errRedisDown
   147  	}
   148  	return tb.real.ForwardIfReady(qnames...)
   149  }
   150  
   151  func (tb *TestBroker) DeleteExpiredCompletedTasks(qname string) error {
   152  	tb.mu.Lock()
   153  	defer tb.mu.Unlock()
   154  	if tb.sleeping {
   155  		return errRedisDown
   156  	}
   157  	return tb.real.DeleteExpiredCompletedTasks(qname)
   158  }
   159  
   160  func (tb *TestBroker) ListLeaseExpired(cutoff time.Time, qnames ...string) ([]*base.TaskMessage, error) {
   161  	tb.mu.Lock()
   162  	defer tb.mu.Unlock()
   163  	if tb.sleeping {
   164  		return nil, errRedisDown
   165  	}
   166  	return tb.real.ListLeaseExpired(cutoff, qnames...)
   167  }
   168  
   169  func (tb *TestBroker) ExtendLease(qname string, ids ...string) (time.Time, error) {
   170  	tb.mu.Lock()
   171  	defer tb.mu.Unlock()
   172  	if tb.sleeping {
   173  		return time.Time{}, errRedisDown
   174  	}
   175  	return tb.real.ExtendLease(qname, ids...)
   176  }
   177  
   178  func (tb *TestBroker) WriteServerState(info *base.ServerInfo, workers []*base.WorkerInfo, ttl time.Duration) error {
   179  	tb.mu.Lock()
   180  	defer tb.mu.Unlock()
   181  	if tb.sleeping {
   182  		return errRedisDown
   183  	}
   184  	return tb.real.WriteServerState(info, workers, ttl)
   185  }
   186  
   187  func (tb *TestBroker) ClearServerState(host string, pid int, serverID string) error {
   188  	tb.mu.Lock()
   189  	defer tb.mu.Unlock()
   190  	if tb.sleeping {
   191  		return errRedisDown
   192  	}
   193  	return tb.real.ClearServerState(host, pid, serverID)
   194  }
   195  
   196  func (tb *TestBroker) CancelationPubSub() (*redis.PubSub, error) {
   197  	tb.mu.Lock()
   198  	defer tb.mu.Unlock()
   199  	if tb.sleeping {
   200  		return nil, errRedisDown
   201  	}
   202  	return tb.real.CancelationPubSub()
   203  }
   204  
   205  func (tb *TestBroker) PublishCancelation(id string) error {
   206  	tb.mu.Lock()
   207  	defer tb.mu.Unlock()
   208  	if tb.sleeping {
   209  		return errRedisDown
   210  	}
   211  	return tb.real.PublishCancelation(id)
   212  }
   213  
   214  func (tb *TestBroker) WriteResult(qname, id string, data []byte) (int, error) {
   215  	tb.mu.Lock()
   216  	defer tb.mu.Unlock()
   217  	if tb.sleeping {
   218  		return 0, errRedisDown
   219  	}
   220  	return tb.real.WriteResult(qname, id, data)
   221  }
   222  
   223  func (tb *TestBroker) Ping() error {
   224  	tb.mu.Lock()
   225  	defer tb.mu.Unlock()
   226  	if tb.sleeping {
   227  		return errRedisDown
   228  	}
   229  	return tb.real.Ping()
   230  }
   231  
   232  func (tb *TestBroker) Close() error {
   233  	tb.mu.Lock()
   234  	defer tb.mu.Unlock()
   235  	if tb.sleeping {
   236  		return errRedisDown
   237  	}
   238  	return tb.real.Close()
   239  }
   240  
   241  func (tb *TestBroker) AddToGroup(ctx context.Context, msg *base.TaskMessage, gname string) error {
   242  	tb.mu.Lock()
   243  	defer tb.mu.Unlock()
   244  	if tb.sleeping {
   245  		return errRedisDown
   246  	}
   247  	return tb.real.AddToGroup(ctx, msg, gname)
   248  }
   249  
   250  func (tb *TestBroker) AddToGroupUnique(ctx context.Context,
   251  	msg *base.TaskMessage, uniqueKey, gname string, ttl time.Duration) error {
   252  	tb.mu.Lock()
   253  	defer tb.mu.Unlock()
   254  	if tb.sleeping {
   255  		return errRedisDown
   256  	}
   257  	return tb.real.AddToGroupUnique(ctx, msg, uniqueKey, gname, ttl)
   258  }
   259  
   260  func (tb *TestBroker) ListGroups(qname string) ([]string, error) {
   261  	tb.mu.Lock()
   262  	defer tb.mu.Unlock()
   263  	if tb.sleeping {
   264  		return nil, errRedisDown
   265  	}
   266  	return tb.real.ListGroups(qname)
   267  }
   268  
   269  func (tb *TestBroker) AggregationCheck(qname, gname string, t time.Time,
   270  	gracePeriod, maxDelay time.Duration, maxSize int) (aggregationSetID string, err error) {
   271  	tb.mu.Lock()
   272  	defer tb.mu.Unlock()
   273  	if tb.sleeping {
   274  		return "", errRedisDown
   275  	}
   276  	return tb.real.AggregationCheck(qname, gname, t, gracePeriod, maxDelay, maxSize)
   277  }
   278  
   279  func (tb *TestBroker) ReadAggregationSet(qname, gname,
   280  	aggregationSetID string) ([]*base.TaskMessage, time.Time, error) {
   281  	tb.mu.Lock()
   282  	defer tb.mu.Unlock()
   283  	if tb.sleeping {
   284  		return nil, time.Time{}, errRedisDown
   285  	}
   286  	return tb.real.ReadAggregationSet(qname, gname, aggregationSetID)
   287  }
   288  
   289  func (tb *TestBroker) DeleteAggregationSet(ctx context.Context, qname, gname, aggregationSetID string) error {
   290  	tb.mu.Lock()
   291  	defer tb.mu.Unlock()
   292  	if tb.sleeping {
   293  		return errRedisDown
   294  	}
   295  	return tb.real.DeleteAggregationSet(ctx, qname, gname, aggregationSetID)
   296  }
   297  
   298  func (tb *TestBroker) ReclaimStaleAggregationSets(qname string) error {
   299  	tb.mu.Lock()
   300  	defer tb.mu.Unlock()
   301  	if tb.sleeping {
   302  		return errRedisDown
   303  	}
   304  	return tb.real.ReclaimStaleAggregationSets(qname)
   305  }