From d8933e1ea576319250ffff88cfb9c3b9774f1a20 Mon Sep 17 00:00:00 2001
From: zeroZshadow <martijn_gerkes@hotmail.com>
Date: Fri, 10 Jun 2016 20:33:36 +0200
Subject: [PATCH] Added tests for MarshalJSON and UnmarshalJSON Reduced
 allocation of MarshalJSON from 65 to 40 B/op

---
 snowflake.go      | 17 ++++++++++-------
 snowflake_test.go | 35 ++++++++++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/snowflake.go b/snowflake.go
index fa6b034..0327a6d 100644
--- a/snowflake.go
+++ b/snowflake.go
@@ -26,9 +26,10 @@ var Epoch int64 = 1288834974657
 // node
 type Node struct {
 	sync.Mutex
-	time int64
-	node int64
-	step int64
+	time     int64
+	node     int64
+	step     int64
+	timebits int64
 }
 
 // An ID is a custom type used for a snowflake ID.  This is used so we can
@@ -51,7 +52,7 @@ func NewNode(node int64) (*Node, error) {
 }
 
 // Generate creates and returns a unique snowflake ID
-func (n *Node) Generate() (ID, error) {
+func (n *Node) Generate() ID {
 
 	n.Lock()
 	defer n.Unlock()
@@ -75,7 +76,7 @@ func (n *Node) Generate() (ID, error) {
 	return ID((now-Epoch)<<timeShift |
 		(n.node << nodeShift) |
 		(n.step),
-	), nil
+	)
 }
 
 // Int64 returns an int64 of the snowflake ID
@@ -125,12 +126,14 @@ func (f ID) Step() int64 {
 
 // MarshalJSON returns a json byte array string of the snowflake ID.
 func (f ID) MarshalJSON() ([]byte, error) {
-	return []byte(`"` + f.String() + `"`), nil
+	buff := []byte("\"")
+	buff = strconv.AppendInt(buff, int64(f), 10)
+	buff = append(buff, '"')
+	return buff, nil
 }
 
 // UnmarshalJSON converts a json byte array of a snowflake ID into an ID type.
 func (f *ID) UnmarshalJSON(b []byte) error {
-
 	i, err := strconv.ParseInt(string(b[1:len(b)-1]), 10, 64)
 	if err != nil {
 		return err
diff --git a/snowflake_test.go b/snowflake_test.go
index dedd8a8..93c459e 100644
--- a/snowflake_test.go
+++ b/snowflake_test.go
@@ -2,6 +2,35 @@ package snowflake
 
 import "testing"
 
+func TestMarshalJSON(t *testing.T) {
+	id := ID(13587)
+	expected := "\"13587\""
+
+	bytes, err := id.MarshalJSON()
+	if err != nil {
+		t.Error("Unexpected error during MarshalJSON")
+	}
+
+	if string(bytes) != expected {
+		t.Errorf("Got %s, expected %s", string(bytes), expected)
+	}
+}
+
+func TestUnmarshalJSON(t *testing.T) {
+	strID := "\"13587\""
+	expected := ID(13587)
+
+	var id ID
+	err := id.UnmarshalJSON([]byte(strID))
+	if err != nil {
+		t.Error("Unexpected error during UnmarshalJSON")
+	}
+
+	if id != expected {
+		t.Errorf("Got %d, expected %d", id, expected)
+	}
+}
+
 func BenchmarkGenerate(b *testing.B) {
 
 	node, _ := NewNode(1)
@@ -10,14 +39,14 @@ func BenchmarkGenerate(b *testing.B) {
 
 	b.ResetTimer()
 	for n := 0; n < b.N; n++ {
-		_, _ = node.Generate()
+		_ = node.Generate()
 	}
 }
 
 func BenchmarkUnmarshal(b *testing.B) {
 
 	node, _ := NewNode(1)
-	id, _ := node.Generate()
+	id := node.Generate()
 	var id2 ID
 
 	b.ReportAllocs()
@@ -31,7 +60,7 @@ func BenchmarkUnmarshal(b *testing.B) {
 func BenchmarkMarshal(b *testing.B) {
 
 	node, _ := NewNode(1)
-	id, _ := node.Generate()
+	id := node.Generate()
 
 	b.ReportAllocs()