Repository: nsqio/nsq
File: nsq/writer_test.go (Go)

Code before:
package nsq

import (
	"errors"
	"log"
	"strconv"
	"testing"
	"time"
)

type ReaderHandler struct {
	t              *testing.T
	q              *Reader
	messagesGood   int
	messagesFailed int
}

func (h *ReaderHandler) LogFailedMessage(message *Message) {
	h.messagesFailed++
	h.q.Stop()
}

func (h *ReaderHandler) HandleMessage(message *Message) error {
	msg := string(message.Body)
	if msg == "bad_test_case" {
		return errors.New("fail this message")
	}
	if msg != "multipublish_test_case" && msg != "publish_test_case" {
		h.t.Error("message 'action' was not correct:", msg)
	}
	h.messagesGood++
	return nil
}
func TestConnection(t *testing.T) {
	addr := "127.0.0.1:4150"
	w := NewWriter()
	err := w.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	w.Stop()
	cmd := Publish("write_test", []byte("failed test"))
	_, _, err = w.Write(cmd)
	if err.Error() != "not connected" {
		t.Fatalf("should not be able to write")
	}
	err = w.ConnectToNSQ(addr)
	if err.Error() != "writer stopped" {
		t.Fatalf("should not be able to connect to NSQ after Stop()")
	}
}

func TestWriteCommand(t *testing.T) {
	addr := "127.0.0.1:4150"
	topicName := "write_test" + strconv.Itoa(int(time.Now().Unix()))
	w := NewWriter()
	err := w.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	msg_count := 10
	for i := 0; i < msg_count; i++ {
		cmd := Publish(topicName, []byte("publish_test_case"))
		frameType, data, err := w.Write(cmd)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
		cmd = Publish(topicName+"a", []byte("publish_test_case"))
		frameType, data, err = w.Write(cmd)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
	}
	var test_data [][]byte
	test_data = append(test_data, []byte("multipublish_test_case"))
	test_data = append(test_data, []byte("multipublish_test_case"))
	for i := 0; i < msg_count; i++ {
		cmd, e := MultiPublish(topicName, test_data)
		if e != nil {
			t.Fatalf("mulitpublish error %s", e.Error())
		}
		frameType, data, err := w.Write(cmd)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
		cmd, e = MultiPublish(topicName+"a", test_data)
		if e != nil {
			t.Fatalf("mulitpublish error %s", e.Error())
		}
		frameType, data, err = w.Write(cmd)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
	}
	cmd := Publish(topicName, []byte("bad_test_case"))
	frameType, data, err := w.Write(cmd)
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	cmd = Publish(topicName+"a", []byte("bad_test_case"))
	frameType, data, err = w.Write(cmd)
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	ReadMessage(topicName, t, msg_count*3)
	ReadMessage(topicName+"a", t, msg_count*3)
}

func TestHeartbeat(t *testing.T) {
	addr := "127.0.0.1:4150"
	topicName := "write_heartbeat_test" + strconv.Itoa(int(time.Now().Unix()))
	w := NewWriter()
	err := w.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	ci := make(map[string]interface{})
	ci["heartbeat_interval"] = 1
	cmd, err := Identify(ci)
	frameType, data, err := w.Write(cmd)
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	time.Sleep(time.Second * 3)
	msg_count := 10
	for i := 0; i < msg_count; i++ {
		cmd = Publish(topicName, []byte("publish_test_case"))
		frameType, data, err = w.Write(cmd)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
	}
	cmd = Publish(topicName, []byte("bad_test_case"))
	frameType, data, err = w.Write(cmd)
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	ReadMessage(topicName, t, msg_count)
}

func ReadMessage(topicName string, t *testing.T, msg_count int) {
	addr := "127.0.0.1:4150"
	q, _ := NewReader(topicName, "ch")
	q.DefaultRequeueDelay = 0
	q.SetMaxBackoffDuration(time.Millisecond * 50)

	h := &ReaderHandler{
		t: t,
		q: q,
	}
	q.AddHandler(h)
	err := q.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	<-q.ExitChan
	log.Println(topicName, "got", h.messagesGood, "and sent", msg_count)
	if h.messagesGood != msg_count {
		t.Fatalf("end of test. should have handled a diff number of messages")
	}
	if h.messagesFailed != 1 {
		t.Fatal("failed message not done")
	}
}


Review comment (suggestion):
this is in milliseconds, so set it to something reasonable like `1000` and then lower the sleep below to 1100 milliseconds.  otherwise this test seems fine to me.


Code after:
package nsq

import (
	"errors"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"
	"time"
)

type ReaderHandler struct {
	t              *testing.T
	q              *Reader
	messagesGood   int
	messagesFailed int
}

func (h *ReaderHandler) LogFailedMessage(message *Message) {
	h.messagesFailed++
	h.q.Stop()
}

func (h *ReaderHandler) HandleMessage(message *Message) error {
	msg := string(message.Body)
	if msg == "bad_test_case" {
		return errors.New("fail this message")
	}
	if msg != "multipublish_test_case" && msg != "publish_test_case" {
		h.t.Error("message 'action' was not correct:", msg)
	}
	h.messagesGood++
	return nil
}
func TestConnection(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	addr := "127.0.0.1:4150"
	w := NewWriter(0)
	err := w.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	w.Stop()
	_, _, err = w.Publish("write_test", []byte("failed test"))
	if err.Error() != "not connected" {
		t.Fatalf("should not be able to write")
	}
	err = w.ConnectToNSQ(addr)
	if err.Error() != "writer stopped" {
		t.Fatalf("should not be able to connect to NSQ after Stop()")
	}
}

func TestWriteCommand(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	addr := "127.0.0.1:4150"
	topicName := "write_test" + strconv.Itoa(int(time.Now().Unix()))
	w := NewWriter(0)
	err := w.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	msg_count := 10
	for i := 0; i < msg_count; i++ {
		frameType, data, err := w.Publish(topicName, []byte("publish_test_case"))
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
		frameType, data, err = w.Publish(topicName+"a", []byte("publish_test_case"))
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
	}
	var test_data [][]byte
	test_data = append(test_data, []byte("multipublish_test_case"))
	test_data = append(test_data, []byte("multipublish_test_case"))
	for i := 0; i < msg_count; i++ {
		frameType, data, err := w.MultiPublish(topicName, test_data)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
		frameType, data, err = w.MultiPublish(topicName+"a", test_data)
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
	}
	frameType, data, err := w.Publish(topicName, []byte("bad_test_case"))
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	frameType, data, err = w.Publish(topicName+"a", []byte("bad_test_case"))
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	ReadMessage(topicName, t, msg_count*3)
	ReadMessage(topicName+"a", t, msg_count*3)
}

func TestHeartbeat(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	addr := "127.0.0.1:4150"
	topicName := "write_heartbeat_test" + strconv.Itoa(int(time.Now().Unix()))
	w := NewWriter(1)
	err := w.ConnectToNSQ(addr)
	if err == nil {
		t.Fatalf("heartbeat test failed")
	}
	_, _, err = w.Publish(topicName, []byte("publish_test_case"))
	if err.Error() != "not connected" {
		t.Fatalf(err.Error())
	}
	w = NewWriter(1000)
	err = w.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	time.Sleep(time.Second * 3)
	msg_count := 10
	for i := 0; i < msg_count; i++ {
		frameType, data, err := w.Publish(topicName, []byte("publish_test_case"))
		if err != nil {
			t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
		}
	}
	frameType, data, err := w.Publish(topicName, []byte("bad_test_case"))
	if err != nil {
		t.Fatalf("frametype %d data %s error %s", frameType, string(data), err.Error())
	}
	ReadMessage(topicName, t, msg_count)
}

func ReadMessage(topicName string, t *testing.T, msg_count int) {
	addr := "127.0.0.1:4150"
	q, _ := NewReader(topicName, "ch")
	q.DefaultRequeueDelay = 0
	q.SetMaxBackoffDuration(time.Millisecond * 50)

	h := &ReaderHandler{
		t: t,
		q: q,
	}
	q.AddHandler(h)
	err := q.ConnectToNSQ(addr)
	if err != nil {
		t.Fatalf(err.Error())
	}
	<-q.ExitChan
	log.Println(topicName, "got", h.messagesGood, "and sent", msg_count)
	if h.messagesGood != msg_count {
		t.Fatalf("end of test. should have handled a diff number of messages")
	}
	if h.messagesFailed != 1 {
		t.Fatal("failed message not done")
	}
}