From 6f08d967ba1e9f3095fb86b39db11ce0994c722a Mon Sep 17 00:00:00 2001
From: anthdm <anthony@academiclabs.co>
Date: Mon, 29 Jan 2018 08:42:49 +0100
Subject: [PATCH] added basic checksum test. + fix travis build

---
 README.md                   |  6 ++++
 pkg/network/message.go      |  6 +++-
 pkg/network/message_test.go | 57 ++++++++-----------------------------
 3 files changed, 23 insertions(+), 46 deletions(-)

diff --git a/README.md b/README.md
index 23a4b53a7..47817401e 100644
--- a/README.md
+++ b/README.md
@@ -10,6 +10,12 @@
     Node and SDK for the <b>NEO</b> blockchain written in the <b>Go</b> language.
 </p>
 
+<p align="center">
+  <a href="https://travis-ci.org/anthdm/neo-go">
+    <img src="https://travis-ci.org/anthdm/neo-go.svg?branch=master">
+  </a>
+</p>
+
 # Overview
 ### Long term project goals
 Full port of the original C# [NEO project](https://github.com/neo-project). A complete toolkit for the NEO blockchain.
diff --git a/pkg/network/message.go b/pkg/network/message.go
index 55d36d927..eeec5a59c 100644
--- a/pkg/network/message.go
+++ b/pkg/network/message.go
@@ -17,6 +17,10 @@ const (
 	cmdSize        = 12
 )
 
+var (
+	errChecksumMismatch = errors.New("checksum mismatch")
+)
+
 // NetMode type that is compatible with netModes below.
 type NetMode uint32
 
@@ -166,7 +170,7 @@ func (m *Message) decodePayload(r io.Reader) error {
 
 	// Compare the checksum of the payload.
 	if !compareChecksum(m.Checksum, pbuf) {
-		return errors.New("checksum mismatch error")
+		return errChecksumMismatch
 	}
 
 	r = bytes.NewReader(pbuf)
diff --git a/pkg/network/message_test.go b/pkg/network/message_test.go
index 35ee154a5..ae443b149 100644
--- a/pkg/network/message_test.go
+++ b/pkg/network/message_test.go
@@ -8,23 +8,6 @@ import (
 	"github.com/anthdm/neo-go/pkg/network/payload"
 )
 
-// func TestNewMessage(t *testing.T) {
-// 	payload := []byte{}
-// 	m := newMessage(ModeTestNet, cmdVersion, payload)
-
-// 	if have, want := m.Length, uint32(0); want != have {
-// 		t.Errorf("want %d have %d", want, have)
-// 	}
-// 	if have, want := len(m.Command), 12; want != have {
-// 		t.Errorf("want %d have %d", want, have)
-// 	}
-
-// 	sum := sumSHA256(sumSHA256(payload))[:4]
-// 	sumuint32 := binary.LittleEndian.Uint32(sum)
-// 	if have, want := m.Checksum, sumuint32; want != have {
-// 		t.Errorf("want %d have %d", want, have)
-// 	}
-// }
 func TestMessageEncodeDecode(t *testing.T) {
 	m := newMessage(ModeTestNet, cmdVersion, nil)
 
@@ -50,52 +33,36 @@ func TestMessageEncodeDecode(t *testing.T) {
 }
 
 func TestMessageEncodeDecodeWithVersion(t *testing.T) {
-	p := payload.NewVersion(2000, "/neo/", 0, true)
+	p := payload.NewVersion(12227, 2000, "./neo:2.6.0/", 0, true)
 	m := newMessage(ModeTestNet, cmdVersion, p)
 
-	buf := &bytes.Buffer{}
+	buf := new(bytes.Buffer)
 	if err := m.encode(buf); err != nil {
 		t.Error(err)
 	}
-	t.Log(buf.Len())
 
-	m1 := &Message{}
-	if err := m1.decode(buf); err != nil {
+	mDecode := &Message{}
+	if err := mDecode.decode(buf); err != nil {
 		t.Fatal(err)
 	}
-	p1 := m1.Payload.(*payload.Version)
 
-	t.Log(p1)
+	if !reflect.DeepEqual(m, mDecode) {
+		t.Fatalf("expected both messages to be equal %v and %v", m, mDecode)
+	}
 }
 
 func TestMessageInvalidChecksum(t *testing.T) {
-	m := newMessage(ModeTestNet, cmdVersion, nil)
+	p := payload.NewVersion(1111, 3000, "./NEO:2.6.0/", 0, true)
+	m := newMessage(ModeTestNet, cmdVersion, p)
 	m.Checksum = 1337
 
-	buf := &bytes.Buffer{}
+	buf := new(bytes.Buffer)
 	if err := m.encode(buf); err != nil {
 		t.Error(err)
 	}
 
 	md := &Message{}
-	if err := md.decode(buf); err == nil {
-		t.Error("decode should failed with checkum mismatch error")
+	if err := md.decode(buf); err == nil && err != errChecksumMismatch {
+		t.Fatalf("decode should fail with %s", errChecksumMismatch)
 	}
 }
-
-// func TestNewVersionPayload(t *testing.T) {
-// 	ua := "/neo/0.0.1/"
-// 	p := newVersionPayload(3000, ua, 0, true)
-// 	b, err := p.encode()
-// 	if err != nil {
-// 		t.Fatal(err)
-// 	}
-
-// 	pd := &Version{}
-// 	if err := pd.decode(b); err != nil {
-// 		t.Fatal(err)
-// 	}
-// 	if !reflect.DeepEqual(p, pd) {
-// 		t.Errorf("both payloads should be equal: %v != %v", p, pd)
-// 	}
-// }