github.com/thiagoyeds/go-cloud@v0.26.0/pubsub/sub_test.go (about)

     1  // Copyright 2018 The Go Cloud Development Kit Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package pubsub_test
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"gocloud.dev/pubsub"
    21  	"gocloud.dev/pubsub/driver"
    22  )
    23  
    24  // scriptedSub returns batches of messages in a predefined order from
    25  // ReceiveBatch.
    26  type scriptedSub struct {
    27  	driver.Subscription
    28  	// batches contains slices of messages to return from ReceiveBatch, one
    29  	// after the other.
    30  	batches [][]*driver.Message
    31  
    32  	// calls counts how many times ReceiveBatch has been called.
    33  	calls int
    34  
    35  	// closed records if Close was called.
    36  	closed bool
    37  }
    38  
    39  func (s *scriptedSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
    40  	b := s.batches[s.calls]
    41  	s.calls++
    42  	return b, nil
    43  }
    44  
    45  func (s *scriptedSub) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
    46  	return nil
    47  }
    48  
    49  func (*scriptedSub) CanNack() bool { return false }
    50  func (s *scriptedSub) Close() error {
    51  	s.closed = true
    52  	return nil
    53  }
    54  
    55  func TestReceiveWithEmptyBatchReturnedFromDriver(t *testing.T) {
    56  	ctx := context.Background()
    57  	ds := &scriptedSub{
    58  		batches: [][]*driver.Message{
    59  			// First call gets an empty batch.
    60  			{},
    61  			// Second call gets a non-empty batch.
    62  			{&driver.Message{}},
    63  		},
    64  	}
    65  	sub := pubsub.NewSubscription(ds, nil, nil)
    66  	defer sub.Shutdown(ctx)
    67  	m, err := sub.Receive(ctx)
    68  	if err != nil {
    69  		t.Fatal(err)
    70  	}
    71  	m.Ack()
    72  }
    73  
    74  func TestSubscriptionCloseIsCalled(t *testing.T) {
    75  	ctx := context.Background()
    76  	ds := &scriptedSub{}
    77  	sub := pubsub.NewSubscription(ds, nil, nil)
    78  	sub.Shutdown(ctx)
    79  	if !ds.closed {
    80  		t.Error("want Subscription.Close to have been called")
    81  	}
    82  }