From 093548c0724c4aabe4c4d57321fe27b27509eea2 Mon Sep 17 00:00:00 2001
From: Waldemar Quevedo <wally@apcera.com>
Date: Wed, 9 Sep 2015 21:09:14 -0700
Subject: [PATCH] Update max payload test to reflect behavior change in go
 client.

Adds logic to test which mimics previous behavior of client.
---
 test/maxpayload_test.go | 60 ++++++++++++++++++++++++++++++++---------
 1 file changed, 48 insertions(+), 12 deletions(-)

diff --git a/test/maxpayload_test.go b/test/maxpayload_test.go
index f372fb0b..29b027f6 100644
--- a/test/maxpayload_test.go
+++ b/test/maxpayload_test.go
@@ -4,9 +4,9 @@ package test
 
 import (
 	"fmt"
+	"net"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/nats-io/nats"
 )
@@ -15,22 +15,58 @@ func TestMaxPayload(t *testing.T) {
 	srv, opts := RunServerWithConfig("./configs/override.conf")
 	defer srv.Shutdown()
 
-	nc, err := nats.Connect(fmt.Sprintf("nats://%s:%d/", opts.Host, opts.Port))
+	endpoint := fmt.Sprintf("%s:%d", opts.Host, opts.Port)
+	nc, err := nats.Connect(fmt.Sprintf("nats://%s/", endpoint))
 	if err != nil {
 		t.Fatalf("Could not connect to server: %v", err)
 	}
 	defer nc.Close()
 
-	big := sizedBytes(4 * 1024 * 1024)
-	nc.Publish("foo", big)
-	err = nc.FlushTimeout(1 * time.Second)
+	size := 4 * 1024 * 1024
+	big := sizedBytes(size)
+	err = nc.Publish("foo", big)
+
+	if err != nats.ErrMaxPayload {
+		t.Fatalf("Expected a Max Payload error")
+	}
+
+	conn, err := net.DialTimeout("tcp", endpoint, nc.Opts.Timeout)
+	if err != nil {
+		t.Fatalf("Could not make a raw connection to the server: %v", err)
+	}
+	info := make([]byte, 200)
+	_, err = conn.Read(info)
+	if err != nil {
+		t.Fatalf("Expected an info message to be sent by the server: %s", err)
+	}
+
+	pub := fmt.Sprintf("PUB bar %d\r\n", size)
+	conn.Write([]byte(pub))
+	if err != nil {
+		t.Fatalf("Could not publish event to the server: %s", err)
+	}
+
+	errMsg := make([]byte, 35)
+	_, err = conn.Read(errMsg)
+	if err != nil {
+		t.Fatalf("Expected an error message to be sent by the server: %s", err)
+	}
+
+	if strings.Contains(string(errMsg), "Maximum Payload Violation") != true {
+		t.Errorf("Received wrong error message (%v)\n", string(errMsg))
+	}
+
+	// Client proactively omits sending the message so server
+	// does not close the connection.
+	if nc.IsClosed() {
+		t.Errorf("Expected connection to not be closed.")
+	}
+
+	// On the other hand client which did not proactively omitted
+	// publishing the bytes following what is suggested by server
+	// in the info message has its connection closed.
+	_, err = conn.Write(big)
 	if err == nil {
-		t.Fatalf("Expected an error from flush")
-	}
-	if strings.Contains(err.Error(), "Maximum Payload Violation") != true {
-		t.Fatalf("Received wrong error message (%v)\n", err)
-	}
-	if !nc.IsClosed() {
-		t.Fatalf("Expected connection to be closed")
+		t.Errorf("Expected error due to maximum payload transgression.")
 	}
 }