From 248f245f9f86758b78337f4edb693f974d3cc24d Mon Sep 17 00:00:00 2001
From: Terence Tsao <terence@prysmaticlabs.com>
Date: Sun, 13 May 2018 11:28:23 -0700
Subject: [PATCH] sharding/contracts: added regression test

Former-commit-id: 01dbc764a0a616d1ed4c21232d5fc886052c3c74 [formerly 5509cecf26e39948254f55d2d2cca820d3ed5b1b]
Former-commit-id: 3830de45de328b24d3d2c81559fac3a1e9d0a76a
---
 sharding/contracts/sharding_manager_test.go | 31 +++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/sharding/contracts/sharding_manager_test.go b/sharding/contracts/sharding_manager_test.go
index 72a27b98f..25267971b 100644
--- a/sharding/contracts/sharding_manager_test.go
+++ b/sharding/contracts/sharding_manager_test.go
@@ -516,6 +516,37 @@ func TestGetCommitteeWithNonMember(t *testing.T) {
 
 }
 
+// TestGetCommitteeWithinSamePeriod tests notary registers and samples within the same period
+func TestGetCommitteeWithinSamePeriod(t *testing.T) {
+	addr := crypto.PubkeyToAddress(mainKey.PublicKey)
+	backend := backends.NewSimulatedBackend(core.GenesisAlloc{addr: {Balance: accountBalance2000Eth}})
+	txOpts := bind.NewKeyedTransactor(mainKey)
+	txOpts.Value = notaryDeposit
+	_, _, smc, _ := deploySMCContract(backend, mainKey)
+
+	// Notary 0 registers.
+	smc.RegisterNotary(txOpts)
+	backend.Commit()
+
+	notary, _ := smc.NotaryRegistry(&bind.CallOpts{}, addr)
+	numNotaries, _ := smc.NotaryPoolLength(&bind.CallOpts{})
+
+	if !notary.Deposited {
+		t.Errorf("Notary has not registered. Got deposited flag: %v", notary.Deposited)
+	}
+
+	if numNotaries.Cmp(big.NewInt(1)) != 0 {
+		t.Errorf("Incorrect count from notary pool. Want: 1, Got: %v", numNotaries)
+	}
+
+	// Notary 0 samples for itself within the same period after registration
+	sampledAddr, _ := smc.GetNotaryInCommittee(&bind.CallOpts{}, big.NewInt(0))
+
+	if addr != sampledAddr {
+		t.Errorf("Unable to sample notary address within same period of registration, got addr: %v", sampledAddr)
+	}
+}
+
 // TestGetCommitteeAfterDeregisters tests notary tries to be in committee after deregistered.
 func TestGetCommitteeAfterDeregisters(t *testing.T) {
 	const notaryCount = 10