From 191568b9575786b0f48981fb71dfb3f2a7e19859 Mon Sep 17 00:00:00 2001
From: Raul Jordan <rauljordan@college.harvard.edu>
Date: Mon, 7 May 2018 16:51:27 -0400
Subject: [PATCH] sharding: test validate shard id

Former-commit-id: 4d83cbd308332a40187191501836e92d3973367c [formerly d5b64486b85afd0c215399b8de520b2d48e8e773]
Former-commit-id: 4a06d18566b3f76eb67405c0e27e78ec06c2144f
---
 sharding/shard.go      |  5 +++++
 sharding/shard_test.go | 24 +++++++++++-------------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/sharding/shard.go b/sharding/shard.go
index fd550d761..abaac5305 100644
--- a/sharding/shard.go
+++ b/sharding/shard.go
@@ -25,6 +25,11 @@ func MakeShard(shardID *big.Int) *Shard {
 	}
 }
 
+// ShardID gets the shard's identifier.
+func (s *Shard) ShardID() *big.Int {
+	return s.shardID
+}
+
 // ValidateShardID checks if header belongs to shard.
 func (s *Shard) ValidateShardID(h *CollationHeader) error {
 	if s.shardID.Cmp(h.shardID) != 0 {
diff --git a/sharding/shard_test.go b/sharding/shard_test.go
index c67736df6..5116e00d4 100644
--- a/sharding/shard_test.go
+++ b/sharding/shard_test.go
@@ -1,24 +1,22 @@
 package sharding
 
 import (
+	"math/big"
 	"testing"
 )
 
 func TestShard_ValidateShardID(t *testing.T) {
-	tests := []struct {
-		headers []*CollationHeader
-	}{
-		{
-			headers: nil,
-		}, {
-			headers: nil,
-		},
+	header := &CollationHeader{shardID: big.NewInt(4)}
+	shard := MakeShard(big.NewInt(3))
+
+	if err := shard.ValidateShardID(header); err == nil {
+		t.Fatalf("Shard ID validation incorrect. Function should throw error when shardID's do not match. want=%d. got=%d", header.shardID.Int64(), shard.ShardID().Int64())
 	}
 
-	for _, tt := range tests {
-		t.Logf("val: %v", tt.headers)
-		if 0 == 1 {
-			t.Fatalf("Wrong number of transactions. want=%d. got=%d", 5, 3)
-		}
+	header2 := &CollationHeader{shardID: big.NewInt(100)}
+	shard2 := MakeShard(big.NewInt(100))
+
+	if err := shard2.ValidateShardID(header2); err != nil {
+		t.Fatalf("Shard ID validation incorrect. Function should not throw error when shardID's match. want=%d. got=%d", header2.shardID.Int64(), shard2.ShardID().Int64())
 	}
 }