consumer_test.go

package nsq

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type MyTestHandler struct {
	t                *testing.T
	q                *Consumer
	messagesSent     int
	messagesReceived int
	messagesFailed   int
	tag              string
	expectFailed     int
	blockingTime     time.Duration
}

type NsqdlookupdWrapper struct {
	lookupdAddr     string
	metaInfoWrapper metaInfo
	t               *testing.T
	fakeResponse    map[string]lookupResp
}

func NewNsqlookupdWrapper(t *testing.T, addr string, metaInfo metaInfo) *NsqdlookupdWrapper {
	proxy := &NsqdlookupdWrapper{
		lookupdAddr:     addr,
		metaInfoWrapper: metaInfo,
		t:               t,
		fakeResponse:    make(map[string]lookupResp),
	}
	return proxy
}

func (self *NsqdlookupdWrapper) fakeListLookupdWrap(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("{}"))
}

func (self *NsqdlookupdWrapper) fakeListLookupdWrapSupplier(response string) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) { w.Write([]byte(response)) }
}

func (self *NsqdlookupdWrapper) fakeLookupdWrap(w http.ResponseWriter, r *http.Request) {
	reqParams, err := url.ParseQuery(r.URL.RawQuery)
	if err != nil {
		self.t.Fatalf("error parse query %v", r.URL.RawQuery)
	}
	topic := reqParams.Get("topic")

	nodes, _ := self.fakeResponse[topic]
	jsonBytes, err := json.Marshal(&nodes)
	if err != nil {
		self.t.Fatalf("error parse lookup resp")
	}
	w.Header().Add("X-Nsq-Content-Type", "nsq; version=1.0")
	self.t.Logf("resp: %s", jsonBytes)
	w.Write(jsonBytes)
}

func (self *NsqdlookupdWrapper) lookupdWrap(w http.ResponseWriter, r *http.Request) {
	reqParams, err := url.ParseQuery(r.URL.RawQuery)
	if err != nil {
		self.t.Fatalf("error parse query %v", r.URL.RawQuery)
	}
	access := reqParams.Get("access")
	topic := reqParams.Get("topic")
	metainfo := reqParams.Get("metainfo")

	val := r.Header.Get("Accept")
	lookupdUrl := fmt.Sprintf("http://%s/lookup?topic=%s&access=%s&metainfo=%s", self.lookupdAddr, topic, access, metainfo)
	req, err := http.NewRequest("GET", lookupdUrl, nil)
	req.Header.Add("Accept", val)
	fmt.Printf("%v\n", val)
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		self.t.Fatalf("error from upstream lookup %v", err)
	}
	var nodes lookupResp
	json.NewDecoder(resp.Body).Decode(&nodes)
	defer resp.Body.Close()
	nodes.Meta = self.metaInfoWrapper
	jsonBytes, err := json.Marshal(&nodes)
	if err != nil {
		self.t.Fatalf("error parse lookup resp")
	}

	w.Header().Add("X-Nsq-Content-Type", "nsq; version=1.0")
	w.Write(jsonBytes)
}

func ensureFakedLookup(t *testing.T, addr string, meta metaInfo, topicList []string, stopC chan struct{}) (*NsqdlookupdWrapper, []*peerInfo, map[string]*peerInfo) {
	lookupdWrapper := NewNsqlookupdWrapper(t, "127.0.0.1:4161", metaInfo{
		PartitionNum:  1,
		Replica:       1,
		ExtendSupport: true,
	})
	allPeers := make([]*peerInfo, 0)
	for i := 0; i < meta.PartitionNum*meta.Replica; i++ {
		var peer peerInfo
		peer.BroadcastAddress = "127.0.0.1"
		peer.Hostname = "test" + strconv.Itoa(i)
		peer.HTTPPort = 4151 + i*100
		peer.RemoteAddress = "127.0.0.1"
		peer.TCPPort = 4150 + i*100
		peer.Version = "1.0"
		allPeers = append(allPeers, &peer)
	}
	usedPeers := make(map[string]*peerInfo)
	for _, t := range topicList {
		rsp := lookupResp{
			Meta:       meta,
			Partitions: make(map[string]*peerInfo),
		}
		for i := 0; i < meta.PartitionNum; i++ {
			peer := allPeers[i]
			rsp.Producers = append(rsp.Producers, peer)
			rsp.Partitions[strconv.Itoa(i)] = peer
			usedPeers[net.JoinHostPort(peer.BroadcastAddress, strconv.Itoa(peer.TCPPort))] = peer
		}
		lookupdWrapper.fakeResponse[t] = rsp
	}
	go func() {
		srvMux := http.NewServeMux()
		srvMux.HandleFunc("/lookup", lookupdWrapper.fakeLookupdWrap)
		srvMux.HandleFunc("/listlookup", lookupdWrapper.fakeListLookupdWrap)
		l, err := net.Listen("tcp", addr)
		if err != nil {
			t.Fatal(err)
		}
		go func() {
			select {
			case <-stopC:
				l.Close()
			}
		}()
		http.Serve(l, srvMux)
	}()
	<-time.After(time.Second)
	return lookupdWrapper, allPeers, usedPeers
}

var nullLogger = log.New(ioutil.Discard, "", log.LstdFlags)

func (h *MyTestHandler) LogFailedMessage(message *Message) {
	h.messagesFailed++
	h.t.Logf("log fail message: %v, %v", h.messagesReceived, GetNewMessageID(message.ID[:]))
	if h.messagesFailed >= h.expectFailed {
		h.q.Stop()
	}
}

func (h *MyTestHandler) HandleMessage(message *Message) error {
	if string(message.Body) == "TOBEFAILED" {
		h.messagesReceived++
		h.t.Logf("received fail message: %v, %v", h.messagesReceived, GetNewMessageID(message.ID[:]))
		return errors.New("fail this message")
	}

	data := struct {
		Msg string
	}{}

	err := json.Unmarshal(message.Body, &data)
	if err != nil {
		return err
	}

	msg := data.Msg
	if msg != "single" && msg != "double" {
		h.t.Error("message 'action' was not correct: ", msg, data)
	}
	if h.tag != "" {
		if message.ExtVer != uint8(4) || !strings.Contains(string(message.ExtBytes), h.tag) {
			h.t.Error("message received has different tag or ext version: ", h.tag, message.ExtVer, string(message.ExtBytes))
		}
	}
	h.messagesReceived++
	h.t.Logf("received message: %v, %v", h.messagesReceived, GetNewMessageID(message.ID[:]))
	if h.blockingTime > 0 && h.messagesReceived > 1 {
		time.Sleep(h.blockingTime)
	}
	return nil
}

type DecompressTestHandler struct {
	t                *testing.T
	q                *Consumer
	messagesSent     int
	messagesReceived int
	messagesFailed   int
	expectFailed     int
	blockingTime     time.Duration
}

func (h *DecompressTestHandler) LogFailedMessage(message *Message) {
	h.messagesFailed++
	h.t.Logf("log fail message: %v, %v", h.messagesReceived, GetNewMessageID(message.ID[:]))
	if h.messagesFailed >= h.expectFailed {
		h.q.Stop()
	}
}

func (h *DecompressTestHandler) HandleMessage(message *Message) error {
	if string(message.Body) == "TOBEFAILED" {
		h.messagesReceived++
		h.t.Logf("received fail message: %v, %v", h.messagesReceived, GetNewMessageID(message.ID[:]))
		return errors.New("fail this message")
	}

	body := message.Body[:]
	if h.q.config.DisableMessageDecompress {
		ext, err := message.GetJsonExt()
		if err != nil {
			h.t.Error("fail to get json ext header: ", message.ExtVer, string(message.ExtBytes))
		}
		//assert compress header
		codecNumStr := ext.Custom[NSQ_CLIENT_COMPRESS_HEADER_KEY].(string)
		if codecNumStr == "" {
			h.t.Error("fail to get nsq compress key json ext header: ", codecNumStr)
		}
		msgSizeStr := ext.Custom[NSQ_CLIENT_COMPRESS_SIZE_HEADER_KEY].(string)
		if msgSizeStr == "" {
			h.t.Error("fail to get nsq compress key json ext header: ", NSQ_CLIENT_COMPRESS_SIZE_HEADER_KEY)
		}

		//decompress
		codecNum, err := strconv.Atoi(codecNumStr)
		if err != nil {
			h.t.Error("fail to convert nsq compress key json ext header: ", codecNumStr)
		}
		codec, err := GetNSQClientCompressCodeByCodecNo(codecNum)
		if err != nil {
			h.t.Error("fail to get compress codec: ", codecNum)
		}
		msgSize, err := strconv.Atoi(msgSizeStr)
		if err != nil {
			h.t.Error("fail to parse message size: ", msgSizeStr)
		}
		body, err = codec.Decompress(body, msgSize)
		if err != nil {
			h.t.Error("fail to decompress message ", err)
		}
	} else {
		//check there is NO nsq client compress ext keys
		ext, err := message.GetJsonExt()
		if err != nil {
			h.t.Error("fail to get json ext header: ", message.ExtVer, string(message.ExtBytes))
		}
		//assert compress header
		if _, ok := ext.Custom[NSQ_CLIENT_COMPRESS_HEADER_KEY]; ok {
			h.t.Error("NSQ_CLIENT_COMPRESS_HEADER_KEY should NOT exist in decompressed message")
		}
		if _, ok := ext.Custom[NSQ_CLIENT_COMPRESS_SIZE_HEADER_KEY]; ok {
			h.t.Error("NSQ_CLIENT_COMPRESS_SIZE_HEADER_KEY should NOT exist in decompressed message")
		}
	}

	data := struct {
		Msg string
	}{}

	err := json.Unmarshal(body, &data)
	if err != nil {
		return err
	}

	msg := data.Msg
	if !strings.HasPrefix(msg, "this is a message for compress test") {
		h.t.Error("message 'action' was not correct: ", msg, data)
	}

	h.messagesReceived++
	h.t.Logf("received message: %v, %v", h.messagesReceived, GetNewMessageID(message.ID[:]))
	if h.blockingTime > 0 && h.messagesReceived > 1 {
		time.Sleep(h.blockingTime)
	}
	return nil
}

func EnsureTopicWithExt(t *testing.T, port int, topic string, part int, ext bool) {
	endpoint := fmt.Sprintf("127.0.0.1:%d", port)
	conn, err := net.DialTimeout("tcp", endpoint, 3*time.Second)
	if err != nil {
		t.Fatalf(err.Error())
		return
	}
	conn.Write(MagicV2)
	if ext {
		CreateTopicWithExt(topic, part).WriteTo(conn)
	} else {
		CreateTopic(topic, part).WriteTo(conn)
	}
	resp, err := ReadResponse(conn)
	if err != nil {
		t.Fatal(err.Error())
		return
	}
	frameType, data, err := UnpackResponse(resp)
	if err != nil {
		t.Fatal(err.Error())
		return
	}
	if frameType == FrameTypeError {
		t.Fatal(string(data))
	}
	time.Sleep(time.Second)
}

func EnsureTopic(t *testing.T, port int, topic string, part int) {
	endpoint := fmt.Sprintf("127.0.0.1:%d", port)
	conn, err := net.DialTimeout("tcp", endpoint, 3*time.Second)
	if err != nil {
		t.Fatalf(err.Error())
		return
	}
	conn.Write(MagicV2)
	CreateTopic(topic, part).WriteTo(conn)
	resp, err := ReadResponse(conn)
	if err != nil {
		t.Fatal(err.Error())
		return
	}
	frameType, data, err := UnpackResponse(resp)
	if err != nil {
		t.Fatal(err.Error())
		return
	}
	if frameType == FrameTypeError {
		t.Fatal(string(data))
	}
	time.Sleep(time.Second)
}

func SendMessage(t *testing.T, port int, topic string, method string, body []byte) {
	httpclient := &http.Client{}
	endpoint := fmt.Sprintf("http://127.0.0.1:%d/%s?topic=%s", port, method, topic)
	req, err := http.NewRequest("POST", endpoint, bytes.NewBuffer(body))
	resp, err := httpclient.Do(req)
	if err != nil {
		t.Fatalf(err.Error())
		return
	}
	resp.Body.Close()
}

func SendCompressMessage(t *testing.T, port int, topic string, method string, body []byte, codecName string) {
	codec, _ := GetNSQClientCompressCodec(codecName)
	compressed, _ := codec.Compress(body)

	httpclient := &http.Client{}
	jsonStr := fmt.Sprintf("{\"%s\":\"%d\",\"%s\":\"%d\"}", NSQ_CLIENT_COMPRESS_SIZE_HEADER_KEY, len(body), NSQ_CLIENT_COMPRESS_HEADER_KEY, codec.GetCodecNo())
	endpoint := fmt.Sprintf("http://127.0.0.1:%d/%s?topic=%s&ext=%s", port, method, topic, url.QueryEscape(jsonStr))
	req, err := http.NewRequest("POST", endpoint, bytes.NewBuffer(compressed))
	resp, err := httpclient.Do(req)
	if err != nil {
		t.Fatalf(err.Error())
		return
	}
	resp.Body.Close()
}

func SendTagedMessage(t *testing.T, port int, topic string, method string, body []byte, tag string) {
	httpclient := &http.Client{}
	jsonStr := fmt.Sprintf("{\"##client_dispatch_tag\":\"%s\"}", tag)
	endpoint := fmt.Sprintf("http://127.0.0.1:%d/%s?topic=%s&ext=%s", port, method, topic, url.QueryEscape(jsonStr))
	req, err := http.NewRequest("POST", endpoint, bytes.NewBuffer(body))
	resp, err := httpclient.Do(req)
	if err != nil {
		t.Fatalf(err.Error())
		return
	}
	resp.Body.Close()
}

func TestConsumer(t *testing.T) {
	consumerTest(t, nil)
}

func TestConsumerTLS(t *testing.T) {
	consumerTest(t, func(c *Config) {
		c.TlsV1 = true
		c.TlsConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	})
}

func TestConsumerDeflate(t *testing.T) {
	consumerTest(t, func(c *Config) {
		c.Deflate = true
	})
}

func TestConsumerSnappy(t *testing.T) {
	consumerTest(t, func(c *Config) {
		c.Snappy = true
	})
}

func TestConsumerTLSDeflate(t *testing.T) {
	consumerTest(t, func(c *Config) {
		c.TlsV1 = true
		c.TlsConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
		c.Deflate = true
	})
}

func TestConsumerTLSSnappy(t *testing.T) {
	consumerTest(t, func(c *Config) {
		c.TlsV1 = true
		c.TlsConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
		c.Snappy = true
	})
}

func TestConsumerTLSClientCert(t *testing.T) {
	envDl := os.Getenv("NSQ_DOWNLOAD")
	if strings.HasPrefix(envDl, "nsq-0.2.24") || strings.HasPrefix(envDl, "nsq-0.2.27") {
		t.Log("skipping due to older nsqd")
		return
	}
	cert, _ := tls.LoadX509KeyPair("./test/client.pem", "./test/client.key")
	consumerTest(t, func(c *Config) {
		c.TlsV1 = true
		c.TlsConfig = &tls.Config{
			Certificates:       []tls.Certificate{cert},
			InsecureSkipVerify: true,
		}
	})
}

func TestConsumerTLSClientCertViaSet(t *testing.T) {
	envDl := os.Getenv("NSQ_DOWNLOAD")
	if strings.HasPrefix(envDl, "nsq-0.2.24") || strings.HasPrefix(envDl, "nsq-0.2.27") {
		t.Log("skipping due to older nsqd")
		return
	}
	consumerTest(t, func(c *Config) {
		c.Set("tls_v1", true)
		c.Set("tls_cert", "./test/client.pem")
		c.Set("tls_key", "./test/client.key")
		c.Set("tls_insecure_skip_verify", true)
	})
}

func TestConsumerDiscoveryLookupd(t *testing.T) {
	// TODO:
}

func TestConsumerSubToNotLeader(t *testing.T) {
	// TODO:
}

func TestConsumerSubToExtTopic(t *testing.T) {
	consumerTagTest(t, func(c *Config) {
		c.Set("desired_tag", "tagTest123")
	})
}

func TestConsumerSubToExtTopicLookupd(t *testing.T) {
	consumerTagTestLookupd(t, func(c *Config) {
		c.Set("desired_tag", "tagTest123")
	})
}

func TestConsumerSubToExtTopicWithCompressedDisableMSGLookupd(t *testing.T) {
	consumerCompressTestLookupd(t, func(c *Config) {
		c.Set("disable_message_decompress", "true")
	})
}

func TestConsumerSubToExtTopicWithCompressedMSGLookupd(t *testing.T) {
	consumerCompressTestLookupd(t, func(c *Config) {
		c.Set("disable_message_decompress", "false")
	})
}

func consumerCompressTestLookupd(t *testing.T, cb func(c *Config)) {
	lookupdAddrWrapper := "127.0.0.1:4162"
	stopCh := make(chan struct{})
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		lookupdWrapper := NewNsqlookupdWrapper(t, "127.0.0.1:4161", metaInfo{
			PartitionNum:  1,
			Replica:       1,
			ExtendSupport: true,
		})

		srvMux := http.NewServeMux()
		srvMux.HandleFunc("/lookup", lookupdWrapper.lookupdWrap)
		//http.ListenAndServe(lookupdAddrWrapper, nil)
		l, err := net.Listen("tcp", lookupdAddrWrapper)
		if err != nil {
			t.Fatal(err)
		}
		go func() {
			select {
			case <-stopCh:
				l.Close()
			}
		}()
		http.Serve(l, srvMux)
	}()
	<-time.After(2 * time.Second)
	fmt.Printf("nsqlookupd starts.")

	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")
	// so that the test can simulate reaching max requeues and a call to LogFailedMessage
	config.DefaultRequeueDelay = 0
	config.MaxAttempts = 5
	// so that the test wont timeout from backing off
	config.MaxBackoffDuration = time.Millisecond * 50
	if cb != nil {
		cb(config)
	}
	config.LookupdPollInterval = time.Second
	//rest desired tag to default
	config.DesiredTag = ""
	topicName := "rdr_compress_test_lookupd"

	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	EnsureTopicWithExt(t, 4150, topicName, 0, true)

	q, _ := NewConsumer(topicName, "ch", config)
	q.SetConsumeExt(true)
	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &DecompressTestHandler{
		t:            t,
		q:            q,
		expectFailed: 0,
	}
	q.AddHandler(h)
	q.ConnectToNSQD("127.0.0.1:4150", 0)

	SendCompressMessage(t, 4151, topicName, "pub_ext", []byte(`{"msg":"this is a message for compress test. XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"}`), CompressDecodec_LZ4.CodecName)
	SendMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"))

	h.messagesSent = 2
	h.expectFailed = 1

	select {
	case <-time.After(time.Second * 20):
		t.Errorf("compressed consumer should stop after timeout")
		q.Stop()
	case <-q.StopChan:
	}

	if h.messagesReceived != h.messagesSent+h.messagesFailed*int(config.MaxAttempts-1) {
		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
	}
	if h.messagesFailed != h.expectFailed {
		t.Fatal("failed message not done")
	}
	close(stopCh)
	wg.Wait()
}

func consumerTagTestLookupd(t *testing.T, cb func(c *Config)) {
	lookupdAddrWrapper := "127.0.0.1:4162"
	stopCh := make(chan struct{})
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		lookupdWrapper := NewNsqlookupdWrapper(t, "127.0.0.1:4161", metaInfo{
			PartitionNum:  1,
			Replica:       1,
			ExtendSupport: true,
		})

		srvMux := http.NewServeMux()
		srvMux.HandleFunc("/lookup", lookupdWrapper.lookupdWrap)
		//http.ListenAndServe(lookupdAddrWrapper, nil)
		l, err := net.Listen("tcp", lookupdAddrWrapper)
		if err != nil {
			t.Fatal(err)
		}
		go func() {
			select {
			case <-stopCh:
				l.Close()
			}
		}()
		http.Serve(l, srvMux)
	}()
	<-time.After(2 * time.Second)
	fmt.Printf("nsqlookupd starts.")

	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")
	// so that the test can simulate reaching max requeues and a call to LogFailedMessage
	config.DefaultRequeueDelay = 0
	config.MaxAttempts = 5
	// so that the test wont timeout from backing off
	config.MaxBackoffDuration = time.Millisecond * 50
	if cb != nil {
		cb(config)
	}
	config.LookupdPollInterval = time.Second
	tag := config.DesiredTag
	//rest desired tag to default
	config.DesiredTag = ""
	topicName := "rdr_tag_test_lookupd"
	if config.Deflate {
		topicName = topicName + "_deflate"
	} else if config.Snappy {
		topicName = topicName + "_snappy"
	}
	if config.TlsV1 {
		topicName = topicName + "_tls"
	}
	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	EnsureTopicWithExt(t, 4150, topicName, 0, true)

	q, _ := NewConsumer(topicName, "ch", config)
	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &MyTestHandler{
		t:            t,
		q:            q,
		expectFailed: 2,
	}
	q.AddHandler(h)
	q.ConnectToNSQLookupd(lookupdAddrWrapper)

	SendTagedMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`), tag)
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendTagedMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`), tag)
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendTagedMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`), tag)
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendTagedMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"), tag)
	SendMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"))
	h.messagesSent = 8

	select {
	case <-time.After(time.Second * 20):
		t.Errorf("tag consumer should stop after timeout")
		q.Stop()
	case <-q.StopChan:
	}

	if h.messagesReceived != h.messagesSent+h.messagesFailed*int(config.MaxAttempts-1) {
		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
	}
	if h.messagesFailed != h.expectFailed {
		t.Fatal("failed message not done")
	}
	close(stopCh)
	wg.Wait()
}

func consumerTagTest(t *testing.T, cb func(c *Config)) {
	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")
	// so that the test can simulate reaching max requeues and a call to LogFailedMessage
	config.DefaultRequeueDelay = 0
	// so that the test wont timeout from backing off
	config.MaxBackoffDuration = time.Millisecond * 50
	config.MaxAttempts = 5
	if cb != nil {
		cb(config)
	}
	tag := config.DesiredTag
	topicName := "rdr_tag_test"
	if config.Deflate {
		topicName = topicName + "_deflate"
	} else if config.Snappy {
		topicName = topicName + "_snappy"
	}
	if config.TlsV1 {
		topicName = topicName + "_tls"
	}
	config.LookupdPollInterval = time.Second

	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	EnsureTopicWithExt(t, 4150, topicName, 0, true)

	q, _ := NewConsumer(topicName, "ch", config)
	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &MyTestHandler{
		t:            t,
		q:            q,
		tag:          tag,
		expectFailed: 1,
	}
	q.AddHandler(h)

	addr := "127.0.0.1:4150"
	//test with invalid tag
	q.SetConsumeExt(true)
	err := q.ConnectToNSQD(addr, 2)
	if err == nil {
		time.Sleep(time.Second)
		// should call on error
		if len(q.connections) > 0 || len(q.nsqdTCPAddrs) > 0 {
			t.Fatal("should disconnect from the NSQ with not exist partition !!!")
		}
	}

	err = q.ConnectToNSQD(addr, 0)
	if err != nil {
		t.Fatal(err)
	}

	//remove desired tag in new consumer
	config.DesiredTag = ""
	qn, _ := NewConsumer(topicName, "ch", config)
	qn.SetConsumeExt(true)
	qn.SetLogger(newTestLogger(t), LogLevelDebug)

	hn := &MyTestHandler{
		t:            t,
		q:            qn,
		expectFailed: 1,
	}
	qn.AddHandler(hn)
	err = qn.ConnectToNSQD(addr, 0)
	if err != nil {
		t.Fatal(err)
	}

	SendTagedMessage(t, 4151, topicName, "pub_ext", []byte(`{"msg":"single"}`), tag)
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendTagedMessage(t, 4151, topicName, "pub_ext", []byte(`{"msg":"single"}`), tag)
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendTagedMessage(t, 4151, topicName, "pub_ext", []byte(`{"msg":"single"}`), tag)
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendTagedMessage(t, 4151, topicName, "pub_ext", []byte("TOBEFAILED"), tag)
	SendMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"))
	h.messagesSent = 4
	hn.messagesSent = 4

	stats := q.Stats()
	if stats.Connections == 0 {
		t.Fatal("stats report 0 connections (should be > 0)")
	}

	err = q.ConnectToNSQD(addr, 0)
	if err == nil {
		t.Fatal("should not be able to connect to the same NSQ twice")
	}

	conn := q.conns()[0]
	if !strings.HasPrefix(conn.conn.LocalAddr().String(), laddr) {
		t.Fatal("connection should be bound to the specified address:", conn.conn.LocalAddr())
	}

	err = q.DisconnectFromNSQD("1.2.3.4:4150", "0")
	if err == nil {
		t.Fatal("should not be able to disconnect from an unknown nsqd")
	}

	err = q.ConnectToNSQD("1.2.3.4:4150", 0)
	if err == nil {
		t.Fatal("should not be able to connect to non-existent nsqd")
	}

	err = q.DisconnectFromNSQD("1.2.3.4:4150", "0")
	if err != nil {
		t.Fatal("should be able to disconnect from an nsqd - " + err.Error())
	}

	select {
	case <-time.After(time.Second * 20):
		t.Errorf("tag consumer should stop after timeout")
		q.Stop()
	case <-q.StopChan:
	}

	select {
	case <-time.After(time.Second * 20):
		t.Errorf("tag consumer should stop after timeout")
		qn.Stop()
	case <-qn.StopChan:
	}

	stats = q.Stats()
	if stats.Connections != 0 {
		t.Fatalf("stats report %d active connections (should be 0), %v", stats.Connections, q.conns())
	}

	stats = q.Stats()
	if stats.MessagesReceived != uint64(h.messagesReceived+h.messagesFailed) {
		t.Fatalf("stats report %d messages received (should be %d)",
			stats.MessagesReceived,
			h.messagesReceived+h.messagesFailed)
	}

	if h.messagesReceived != h.messagesSent+h.messagesFailed*int(config.MaxAttempts-1) {
		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
	}
	if h.messagesFailed != h.expectFailed {
		t.Fatal("failed message not done")
	}

	stats = qn.Stats()
	if stats.Connections != 0 {
		t.Fatalf("stats report %d active connections (should be 0)", stats.Connections)
	}

	stats = qn.Stats()
	if stats.MessagesReceived != uint64(hn.messagesReceived+hn.messagesFailed) {
		t.Fatalf("stats report %d messages received (should be %d)",
			stats.MessagesReceived,
			hn.messagesReceived+hn.messagesFailed)
	}

	if hn.messagesReceived != hn.messagesSent+hn.messagesFailed*(int(config.MaxAttempts)-1) {
		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", hn.messagesReceived, hn.messagesSent)
	}
	if hn.messagesFailed != hn.expectFailed {
		t.Fatal("failed message not done")
	}
}

func consumerTest(t *testing.T, cb func(c *Config)) {
	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")
	// so that the test can simulate reaching max requeues and a call to LogFailedMessage
	config.DefaultRequeueDelay = 0
	// so that the test wont timeout from backing off
	config.MaxBackoffDuration = time.Millisecond * 50
	config.MaxAttempts = 7
	config.LookupdPollInterval = time.Second
	if cb != nil {
		cb(config)
	}
	topicName := "rdr_test"
	if config.Deflate {
		topicName = topicName + "_deflate"
	} else if config.Snappy {
		topicName = topicName + "_snappy"
	}
	if config.TlsV1 {
		topicName = topicName + "_tls"
	}
	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	q, _ := NewConsumer(topicName, "ch", config)
	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &MyTestHandler{
		t:            t,
		q:            q,
		expectFailed: 1,
	}
	q.AddHandler(h)

	EnsureTopic(t, 4150, topicName, 0)

	addr := "127.0.0.1:4150"
	err := q.ConnectToNSQD(addr, 2)
	if err == nil {
		time.Sleep(time.Second)
		// should call on error
		if len(q.connections) > 0 || len(q.nsqdTCPAddrs) > 0 {
			t.Fatal("should disconnect from the NSQ with not exist partition !!!")
		}
	}

	err = q.ConnectToNSQD(addr, 0)
	if err != nil {
		t.Fatal(err)
	}
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendMessage(t, 4151, topicName, "mpub", []byte("{\"msg\":\"double\"}\n{\"msg\":\"double\"}"))
	SendMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"))
	h.messagesSent = 4

	stats := q.Stats()
	if stats.Connections == 0 {
		t.Fatal("stats report 0 connections (should be > 0)")
	}

	err = q.ConnectToNSQD(addr, 0)
	if err == nil {
		t.Fatal("should not be able to connect to the same NSQ twice")
	}

	conn := q.conns()[0]
	if !strings.HasPrefix(conn.conn.LocalAddr().String(), laddr) {
		t.Fatal("connection should be bound to the specified address:", conn.conn.LocalAddr())
	}

	err = q.DisconnectFromNSQD("1.2.3.4:4150", "0")
	if err == nil {
		t.Fatal("should not be able to disconnect from an unknown nsqd")
	}

	err = q.ConnectToNSQD("1.2.3.4:4150", 0)
	if err == nil {
		t.Fatal("should not be able to connect to non-existent nsqd")
	}

	err = q.DisconnectFromNSQD("1.2.3.4:4150", "0")
	if err != nil {
		t.Fatal("should be able to disconnect from an nsqd - " + err.Error())
	}

	select {
	case <-q.StopChan:
	case <-time.After(time.Second * 10):
		t.Errorf("should stop after timeout")
		q.Stop()
	}

	stats = q.Stats()
	if stats.Connections != 0 {
		t.Fatalf("stats report %d active connections (should be 0)", stats.Connections)
	}

	stats = q.Stats()
	if stats.MessagesReceived != uint64(h.messagesReceived+h.messagesFailed) {
		t.Fatalf("stats report %d messages received (should be %d)",
			stats.MessagesReceived,
			h.messagesReceived+h.messagesFailed)
	}

	if h.messagesReceived != h.messagesSent+int(config.MaxAttempts)-1 {
		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
	}
	if h.messagesFailed != h.expectFailed {
		t.Fatal("failed message not done")
	}
}

func TestConsumerBackoffOnDifferentNode(t *testing.T) {
	// test backoff on one node should not affect the other node
}

func TestConsumerBlockingOnHandler(t *testing.T) {
	maxCleanupWaiting = time.Second * 10
	maxStopWaiting = time.Second * 5
	// block handler and test the consumer call r.exit() early before the handler returned
	// block handler and test finish after the conn cleanup
	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")
	// so that the test can simulate reaching max requeues and a call to LogFailedMessage
	config.DefaultRequeueDelay = 0
	// so that the test wont timeout from backing off
	config.MaxBackoffDuration = time.Millisecond * 50
	config.MaxAttempts = 7
	config.LookupdPollInterval = time.Second
	config.MsgTimeout = time.Minute

	topicName := "rdr_blocking_test"
	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	q, _ := NewConsumer(topicName, "ch", config)

	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &MyTestHandler{
		t:            t,
		q:            q,
		expectFailed: 1,
		blockingTime: time.Second * 15,
	}
	q.AddHandler(h)

	EnsureTopic(t, 4150, topicName, 0)

	addr := "127.0.0.1:4150"
	err := q.ConnectToNSQD(addr, 0)
	if err != nil {
		t.Fatal(err)
	}
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendMessage(t, 4151, topicName, "mpub", []byte("{\"msg\":\"double\"}\n{\"msg\":\"double\"}"))
	SendMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"))
	h.messagesSent = 4

	stats := q.Stats()
	if stats.Connections == 0 {
		t.Fatal("stats report 0 connections (should be > 0)")
	}
	time.Sleep(time.Second)
	s := time.Now()
	q.Stop()
	select {
	case <-q.StopChan:
	case <-time.After(h.blockingTime):
		t.Errorf("should stop after timeout")
	}

	stats = q.Stats()
	if stats.Connections == 0 {
		t.Errorf("should have active connections while blocking consumer")
	}
	for {
		stats = q.Stats()
		if stats.Connections == 0 {
			break
		}
		time.Sleep(time.Second)
	}
	cost := time.Since(s)
	if cost > h.blockingTime*time.Duration(h.messagesSent) {
		t.Errorf("should not blocking too long: %s", cost)
	}
	stats = q.Stats()
	if stats.MessagesReceived != uint64(h.messagesReceived+h.messagesFailed) {
		t.Errorf("stats report %d messages received (should be %d+%d)",
			stats.MessagesReceived,
			h.messagesReceived, h.messagesFailed)
	}
	if h.messagesReceived != h.messagesSent {
		t.Errorf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
	}
	// the failed message will not be requeued after stopped, so it will be ignored
	if h.messagesFailed != 0 {
		t.Errorf("failed message not done: %v, %v", h.messagesFailed, h.expectFailed)
	}
}

func TestConsumerBlockingOnHandlerWhileConnClosed(t *testing.T) {
	maxCleanupWaiting = time.Second * 10
	maxStopWaiting = time.Second * 5
	// block handler and test the consumer call r.exit() early before the handler returned
	// block handler and test finish after the conn cleanup
	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")
	// so that the test can simulate reaching max requeues and a call to LogFailedMessage
	config.DefaultRequeueDelay = 0
	// so that the test wont timeout from backing off
	config.MaxBackoffDuration = time.Millisecond * 50
	config.MaxAttempts = 7
	config.LookupdPollInterval = time.Second
	config.MsgTimeout = time.Minute

	topicName := "rdr_blocking_test_connclosed"
	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	q, _ := NewConsumer(topicName, "ch", config)

	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &MyTestHandler{
		t:            t,
		q:            q,
		expectFailed: 1,
		blockingTime: time.Second * 15,
	}
	q.AddHandler(h)

	EnsureTopic(t, 4150, topicName, 0)

	addr := "127.0.0.1:4150"
	err := q.ConnectToNSQD(addr, 0)
	if err != nil {
		t.Fatal(err)
	}
	SendMessage(t, 4151, topicName, "pub", []byte(`{"msg":"single"}`))
	SendMessage(t, 4151, topicName, "mpub", []byte("{\"msg\":\"double\"}\n{\"msg\":\"double\"}"))
	SendMessage(t, 4151, topicName, "pub", []byte("TOBEFAILED"))
	h.messagesSent = 4

	stats := q.Stats()
	if stats.Connections == 0 {
		t.Fatal("stats report 0 connections (should be > 0)")
	}

	time.Sleep(time.Second)
	s := time.Now()
	q.Stop()
	select {
	case <-q.StopChan:
	case <-time.After(h.blockingTime):
		t.Errorf("should stop after timeout")
	}

	stats = q.Stats()
	if stats.Connections != 0 {
		t.Logf("stats report %d active connections", stats.Connections)
	}

	// force close the connection to trigger the cleanup on connection
	for _, c := range q.conns() {
		c.conn.Close()
	}

	for {
		stats = q.Stats()
		if stats.Connections == 0 {
			break
		}
		time.Sleep(time.Second)
		if time.Since(s) > time.Minute {
			t.Errorf("timeout wait")
		}
	}
	cost := time.Since(s)
	if cost > h.blockingTime*time.Duration(h.messagesSent) {
		t.Errorf("should not blocking too long: %s", cost)
	}
	stats = q.Stats()
	if stats.MessagesReceived != uint64(h.messagesSent) {
		t.Errorf("stats report %d messages received (should be %d+%d)",
			stats.MessagesReceived,
			h.messagesReceived, h.messagesFailed)
	}
	if int(stats.MessagesFinished) >= h.messagesSent-1 {
		t.Errorf("end of test. should ignore some number of messages (got %d, sent %d)", stats, h.messagesSent)
	}
	if h.messagesReceived+h.messagesFailed >= h.messagesSent {
		t.Errorf("end of test. should ignore some number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
	}
	// the failed message will not be requeued after stopped, so it will be ignored
	if h.messagesFailed != 0 {
		t.Errorf("failed message should be ignored: %v, %v", h.messagesFailed, h.expectFailed)
	}
	for {
		stats = q.Stats()
		if int(stats.MessagesFinished) == h.messagesSent {
			break
		}
		if atomic.LoadInt32(&q.runningHandlers) == 0 {
			break
		}
		time.Sleep(time.Second)
		if time.Since(s) > h.blockingTime*time.Duration(h.messagesSent) {
			t.Errorf("timeout wait")
			break
		}
	}
}

func TestConsumerConnectFailedShouldCleanConn(t *testing.T) {
	maxCleanupWaiting = time.Second * 10
	maxStopWaiting = time.Second * 5
	// block handler and test the consumer call r.exit() early before the handler returned
	// block handler and test finish after the conn cleanup
	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")

	topicName := "rdr_test_connfailed"
	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	q, _ := NewConsumer(topicName, "ch", config)

	q.SetLogger(newTestLogger(t), LogLevelDebug)

	h := &MyTestHandler{
		t:            t,
		q:            q,
		expectFailed: 1,
		blockingTime: time.Second * 15,
	}
	q.AddHandler(h)

	addr := "127.0.0.1:4150"
	err := q.ConnectToNSQD(addr, 0)
	if err == nil {
		t.Errorf("should connect failed")
	}

	stats := q.Stats()
	if stats.Connections > 0 {
		t.Errorf("stats report connections (should be  0), %v", stats)
	}
	time.Sleep(time.Second)
	q.Stop()
	select {
	case <-q.StopChan:
	case <-time.After(h.blockingTime):
		t.Errorf("should stop after timeout")
	}
	time.Sleep(time.Second * 2)
}

func TestConsumerBackoffWhileChangeMaxInflight(t *testing.T) {
	topicName := "nsq_test"
	config := NewConfig()

	config.MaxInFlight = 80

	q, _ := NewConsumer(topicName, "ch", config)
	q.AddHandler(&testHandler{})
	EnsureTopic(t, 4150, topicName, 0)
	err := q.ConnectToNSQLookupd("127.0.0.1:4161")
	if err != nil {
		t.Fatalf(err.Error())
	}
	time.Sleep(5 * time.Second)
	conns := q.conns()
	if len(conns) == 0 {
		t.Fatalf("fail to get consumer conns")
	}

	try := 3
	for try > 0 {
		q.onConnBackoff(conns[0], true)
		// change to low max inflight
		time.Sleep(time.Minute/4)
		q.ChangeMaxInFlight(20)
		time.Sleep(time.Second * 5)
		for _, con := range q.conns() {
			rdy := con.RDY()
			if rdy <= 0 {
				t.Fatalf("rdy should not be zero %v", rdy)
			}
		}
		try--
	}
}

type testConnDelegate struct {
	OnClosed int32
}

func (d *testConnDelegate) OnResponse(c *Conn, data []byte)       {}
func (d *testConnDelegate) OnError(c *Conn, data []byte)          {}
func (d *testConnDelegate) OnMessage(c *Conn, m *Message)         {}
func (d *testConnDelegate) OnMessageFinished(c *Conn, m *Message) {}
func (d *testConnDelegate) OnMessageRequeued(c *Conn, m *Message) {}
func (d *testConnDelegate) OnBackoff(c *Conn, connOnly bool)      {}
func (d *testConnDelegate) OnContinue(c *Conn)                    {}
func (d *testConnDelegate) OnResume(c *Conn)                      {}
func (d *testConnDelegate) OnIOError(c *Conn, err error)          {}
func (d *testConnDelegate) OnHeartbeat(c *Conn)                   {}
func (d *testConnDelegate) OnClose(c *Conn) {
	atomic.StoreInt32(&d.OnClosed, 1)
}

func TestConnConnectFailedShouldCleanConn(t *testing.T) {
	config := NewConfig()
	laddr := "127.0.0.1"
	// so that the test can simulate binding consumer to specified address
	config.LocalAddr, _ = net.ResolveTCPAddr("tcp", laddr+":0")

	topicName := "rdr_test_connfailed"
	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
	testDelegate := &testConnDelegate{}
	conn := NewConn("127.0.0.1:4150", config, testDelegate)
	conn.consumePart = "0"

	cleanupConnection := func() {
		conn.CloseAll()
	}
	_, err := conn.Connect()
	if err != nil {
		cleanupConnection()
	}
	time.Sleep(time.Second)
	if atomic.LoadInt32(&testDelegate.OnClosed) != 1 {
		t.Errorf("should trigger onclosed after cleanup")
	}
}