From 685e200ee3322ea7bcd8c965ab424410045cef57 Mon Sep 17 00:00:00 2001
From: Leo Chen <leo@harmony.one>
Date: Sat, 27 Jul 2019 16:00:19 -0700
Subject: [PATCH] [test] add unit tests for resharding func

Signed-off-by: Leo Chen <leo@harmony.one>
---
 core/core_test.go                             | 67 +++++++++++++++++
 .../configs/sharding/shardingconfig_test.go   | 72 +++++++++++++++++++
 2 files changed, 139 insertions(+)
 create mode 100644 core/core_test.go
 create mode 100644 internal/configs/sharding/shardingconfig_test.go

diff --git a/core/core_test.go b/core/core_test.go
new file mode 100644
index 000000000..bbee042d4
--- /dev/null
+++ b/core/core_test.go
@@ -0,0 +1,67 @@
+package core
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/harmony-one/harmony/core/types"
+	shardingconfig "github.com/harmony-one/harmony/internal/configs/sharding"
+)
+
+func TestIsEpochBlock(t *testing.T) {
+	block1 := types.NewBlock(&types.Header{Number: big.NewInt(10)}, nil, nil)
+	block2 := types.NewBlock(&types.Header{Number: big.NewInt(0)}, nil, nil)
+	block3 := types.NewBlock(&types.Header{Number: big.NewInt(327680)}, nil, nil)
+	block4 := types.NewBlock(&types.Header{Number: big.NewInt(77)}, nil, nil)
+	block5 := types.NewBlock(&types.Header{Number: big.NewInt(78)}, nil, nil)
+	block6 := types.NewBlock(&types.Header{Number: big.NewInt(188)}, nil, nil)
+	block7 := types.NewBlock(&types.Header{Number: big.NewInt(189)}, nil, nil)
+	tests := []struct {
+		schedule shardingconfig.Schedule
+		block    *types.Block
+		expected bool
+	}{
+		{
+			shardingconfig.MainnetSchedule,
+			block1,
+			false,
+		},
+		{
+			shardingconfig.MainnetSchedule,
+			block2,
+			true,
+		},
+		{
+			shardingconfig.MainnetSchedule,
+			block3,
+			true,
+		},
+		{
+			shardingconfig.TestnetSchedule,
+			block4,
+			false,
+		},
+		{
+			shardingconfig.TestnetSchedule,
+			block5,
+			true,
+		},
+		{
+			shardingconfig.TestnetSchedule,
+			block6,
+			false,
+		},
+		{
+			shardingconfig.TestnetSchedule,
+			block7,
+			true,
+		},
+	}
+	for _, test := range tests {
+		ShardingSchedule = test.schedule
+		r := IsEpochBlock(test.block)
+		if r != test.expected {
+			t.Errorf("expected: %v, got: %v\n", test.expected, r)
+		}
+	}
+}
diff --git a/internal/configs/sharding/shardingconfig_test.go b/internal/configs/sharding/shardingconfig_test.go
new file mode 100644
index 000000000..f1b05ed11
--- /dev/null
+++ b/internal/configs/sharding/shardingconfig_test.go
@@ -0,0 +1,72 @@
+package shardingconfig
+
+import (
+	"math/big"
+	"testing"
+)
+
+func TestMainnetInstanceForEpoch(t *testing.T) {
+	tests := []struct {
+		epoch    *big.Int
+		instance Instance
+	}{
+		{
+			big.NewInt(0),
+			mainnetV0,
+		},
+		{
+			big.NewInt(1),
+			mainnetV1,
+		},
+		{
+			big.NewInt(2),
+			mainnetV1,
+		},
+	}
+
+	for _, test := range tests {
+		in := MainnetSchedule.InstanceForEpoch(test.epoch)
+		if in.NumShards() != test.instance.NumShards() || in.NumNodesPerShard() != test.instance.NumNodesPerShard() {
+			t.Errorf("can't get the right instane for epoch: %v\n", test.epoch)
+		}
+	}
+}
+
+func TestCalcEpochNumber(t *testing.T) {
+	tests := []struct {
+		block uint64
+		epoch *big.Int
+	}{
+		{
+			0,
+			big.NewInt(0),
+		},
+		{
+			1,
+			big.NewInt(0),
+		},
+		{
+			327679,
+			big.NewInt(0),
+		},
+		{
+			327680,
+			big.NewInt(1),
+		},
+		{
+			344064,
+			big.NewInt(2),
+		},
+		{
+			344063,
+			big.NewInt(1),
+		},
+	}
+
+	for _, test := range tests {
+		ep := MainnetSchedule.CalcEpochNumber(test.block)
+		if ep.Cmp(test.epoch) != 0 {
+			t.Errorf("CalcEpochNumber error: got %v, expect %v\n", ep, test.epoch)
+		}
+	}
+}