github.com/Psiphon-Inc/goarista@v0.0.0-20160825065156-d002785f4c67/kafka/producer/producer_test.go (about)

     1  // Copyright (C) 2016  Arista Networks, Inc.
     2  // Use of this source code is governed by the Apache License 2.0
     3  // that can be found in the COPYING file.
     4  
     5  package producer
     6  
     7  import (
     8  	"encoding/json"
     9  	"strings"
    10  	"sync"
    11  	"testing"
    12  
    13  	"github.com/Shopify/sarama"
    14  	"github.com/aristanetworks/goarista/kafka/openconfig"
    15  	pb "github.com/aristanetworks/goarista/openconfig"
    16  	"github.com/aristanetworks/goarista/test"
    17  	"github.com/golang/protobuf/proto"
    18  )
    19  
    20  type mockAsyncProducer struct {
    21  	input     chan *sarama.ProducerMessage
    22  	successes chan *sarama.ProducerMessage
    23  	errors    chan *sarama.ProducerError
    24  }
    25  
    26  func newMockAsyncProducer() *mockAsyncProducer {
    27  	return &mockAsyncProducer{
    28  		input:     make(chan *sarama.ProducerMessage),
    29  		successes: make(chan *sarama.ProducerMessage),
    30  		errors:    make(chan *sarama.ProducerError)}
    31  }
    32  
    33  func (p *mockAsyncProducer) AsyncClose() {
    34  	panic("Not implemented")
    35  }
    36  
    37  func (p *mockAsyncProducer) Close() error {
    38  	close(p.successes)
    39  	close(p.errors)
    40  	return nil
    41  }
    42  
    43  func (p *mockAsyncProducer) Input() chan<- *sarama.ProducerMessage {
    44  	return p.input
    45  }
    46  
    47  func (p *mockAsyncProducer) Successes() <-chan *sarama.ProducerMessage {
    48  	return p.successes
    49  }
    50  
    51  func (p *mockAsyncProducer) Errors() <-chan *sarama.ProducerError {
    52  	return p.errors
    53  }
    54  
    55  func newPath(path string) *pb.Path {
    56  	if path == "" {
    57  		return nil
    58  	}
    59  	return &pb.Path{Element: strings.Split(path, "/")}
    60  }
    61  
    62  func TestKafkaProducer(t *testing.T) {
    63  	mock := newMockAsyncProducer()
    64  	toDB := make(chan proto.Message)
    65  	topic := "occlient"
    66  	systemID := "Foobar"
    67  	toDBProducer := &producer{
    68  		notifsChan:    toDB,
    69  		kafkaProducer: mock,
    70  		topic:         topic,
    71  		key:           sarama.StringEncoder(systemID),
    72  		encoder:       openconfig.ElasticsearchMessageEncoder,
    73  		done:          make(chan struct{}),
    74  		wg:            sync.WaitGroup{},
    75  	}
    76  
    77  	go toDBProducer.Run()
    78  
    79  	response := &pb.SubscribeResponse{
    80  		Response: &pb.SubscribeResponse_Update{
    81  			Update: &pb.Notification{
    82  				Timestamp: 0,
    83  				Prefix:    newPath("/foo/bar"),
    84  				Update:    []*pb.Update{},
    85  			},
    86  		},
    87  	}
    88  	document := map[string]interface{}{
    89  		"": map[string]interface{}{
    90  			"foo": map[string]interface{}{
    91  				"bar": map[string]interface{}{},
    92  			},
    93  		},
    94  	}
    95  
    96  	toDB <- response
    97  
    98  	kafkaMessage := <-mock.input
    99  	if kafkaMessage.Topic != topic {
   100  		t.Errorf("Unexpected Topic: %s, expecting %s", kafkaMessage.Topic, topic)
   101  	}
   102  	key, err := kafkaMessage.Key.Encode()
   103  	if err != nil {
   104  		t.Fatalf("Error encoding key: %s", err)
   105  	}
   106  	if string(key) != systemID {
   107  		t.Errorf("Kafka message didn't have expected key: %s, expecting %s", string(key), systemID)
   108  	}
   109  
   110  	valueBytes, err := kafkaMessage.Value.Encode()
   111  	if err != nil {
   112  		t.Fatalf("Error encoding value: %s", err)
   113  	}
   114  	var result interface{}
   115  	err = json.Unmarshal(valueBytes, &result)
   116  	if err != nil {
   117  		t.Errorf("Error decoding into JSON: %s", err)
   118  	}
   119  	if !test.DeepEqual(document[""], result.(map[string]interface{})[""]) {
   120  		t.Errorf("Protobuf sent from Kafka Producer does not match original.\nOriginal: %v\nNew:%v",
   121  			document, result)
   122  	}
   123  
   124  	toDBProducer.Stop()
   125  }