github.com/adnan-c/fabric_e2e_couchdb@v0.6.1-preview.0.20170228180935-21ce6b23cf91/orderer/kafka/consumer_mock_test.go (about)

     1  /*
     2  Copyright IBM Corp. 2016 All Rights Reserved.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8                   http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package kafka
    18  
    19  import (
    20  	"fmt"
    21  	"testing"
    22  	"time"
    23  
    24  	ab "github.com/hyperledger/fabric/protos/orderer"
    25  	"github.com/hyperledger/fabric/protos/utils"
    26  
    27  	"github.com/Shopify/sarama"
    28  	"github.com/Shopify/sarama/mocks"
    29  )
    30  
    31  type mockConsumerImpl struct {
    32  	consumedOffset int64
    33  	chainPartition ChainPartition
    34  
    35  	parentConsumer         *mocks.Consumer
    36  	chainPartitionManager  *mocks.PartitionConsumer
    37  	chainPartitionConsumer sarama.PartitionConsumer
    38  
    39  	disk         chan *ab.KafkaMessage
    40  	isSetup      chan struct{}
    41  	targetOffset int64
    42  	t            *testing.T
    43  }
    44  
    45  func mockNewConsumer(t *testing.T, cp ChainPartition, offset int64, disk chan *ab.KafkaMessage) (Consumer, error) {
    46  	var err error
    47  	parentConsumer := mocks.NewConsumer(t, nil)
    48  	// NOTE The seek flag seems to be useless here.
    49  	// The mock partition will have its highWatermarkOffset
    50  	// initialized to 0 no matter what. I've opened up an issue
    51  	// in the sarama repo: https://github.com/Shopify/sarama/issues/745
    52  	// Until this is resolved, use the testFillWithBlocks() hack below.
    53  	cpManager := parentConsumer.ExpectConsumePartition(cp.Topic(), cp.Partition(), offset)
    54  	cpConsumer, err := parentConsumer.ConsumePartition(cp.Topic(), cp.Partition(), offset)
    55  	// mockNewConsumer is basically a helper function when testing.
    56  	// Any errors it generates internally, should result in panic
    57  	// and not get propagated further; checking its errors in the
    58  	// calling functions (i.e. the actual tests) increases boilerplate.
    59  	if err != nil {
    60  		t.Fatal("Cannot create mock partition consumer:", err)
    61  	}
    62  	mc := &mockConsumerImpl{
    63  		consumedOffset: 0,
    64  		targetOffset:   offset,
    65  		chainPartition: cp,
    66  
    67  		parentConsumer:         parentConsumer,
    68  		chainPartitionManager:  cpManager,
    69  		chainPartitionConsumer: cpConsumer,
    70  		disk:    disk,
    71  		isSetup: make(chan struct{}),
    72  		t:       t,
    73  	}
    74  	// Stop-gap hack until sarama issue #745 is resolved:
    75  	if mc.targetOffset >= testOldestOffset && mc.targetOffset <= (testNewestOffset-1) {
    76  		mc.testFillWithBlocks(mc.targetOffset - 1) // Prepare the consumer so that the next Recv gives you blob #targetOffset
    77  	} else {
    78  		err = fmt.Errorf("Out of range offset (seek number) given to consumer: %d", offset)
    79  		return mc, err
    80  	}
    81  
    82  	return mc, err
    83  }
    84  
    85  func (mc *mockConsumerImpl) Recv() <-chan *sarama.ConsumerMessage {
    86  	if mc.consumedOffset >= testNewestOffset-1 {
    87  		return nil
    88  	}
    89  
    90  	// This is useful in cases where we want to <-Recv() in a for/select loop in
    91  	// a non-blocking manner. Without the timeout, the Go runtime will always
    92  	// execute the body of the Recv() method. If there in no outgoing message
    93  	// available, it will block while waiting on mc.disk. All the other cases in
    94  	// the original for/select loop then won't be evaluated until we unblock on
    95  	// <-mc.disk (which may never happen).
    96  	select {
    97  	case <-time.After(testTimePadding / 2):
    98  	case outgoingMsg := <-mc.disk:
    99  		mc.consumedOffset++
   100  		mc.chainPartitionManager.YieldMessage(testNewConsumerMessage(mc.chainPartition, mc.consumedOffset, outgoingMsg))
   101  		if mc.consumedOffset == mc.targetOffset-1 {
   102  			close(mc.isSetup) // Hook for callers
   103  		}
   104  		return mc.chainPartitionConsumer.Messages()
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (mc *mockConsumerImpl) Errors() <-chan *sarama.ConsumerError {
   111  	return nil
   112  }
   113  
   114  func (mc *mockConsumerImpl) Close() error {
   115  	if err := mc.chainPartitionManager.Close(); err != nil {
   116  		return err
   117  	}
   118  	return mc.parentConsumer.Close()
   119  }
   120  
   121  func (mc *mockConsumerImpl) testFillWithBlocks(offset int64) {
   122  	for i := int64(1); i <= offset; i++ {
   123  		go func() {
   124  			mc.disk <- newRegularMessage(utils.MarshalOrPanic(newTestEnvelope(fmt.Sprintf("consumer fill-in %d", i))))
   125  		}()
   126  		<-mc.Recv()
   127  	}
   128  	return
   129  }
   130  
   131  func testNewConsumerMessage(cp ChainPartition, offset int64, kafkaMessage *ab.KafkaMessage) *sarama.ConsumerMessage {
   132  	return &sarama.ConsumerMessage{
   133  		Value:     sarama.ByteEncoder(utils.MarshalOrPanic(kafkaMessage)),
   134  		Topic:     cp.Topic(),
   135  		Partition: cp.Partition(),
   136  		Offset:    offset,
   137  	}
   138  }