add test for sharding structure

internal/hmyapi/blockchain.go

@@ -26,6 +26,14 @@ const (
 	defaultFromAddress = "0x0000000000000000000000000000000000000000"
 )
 
+// Constants for sharding structure.
+const (
+	MainNetHTTPPattern = "http://s%d.t.hmny.io:9500"
+	MainNetWSPattern   = "ws://s%d.t.hmny.io:9800"
+	TestNetHTTPPattern = "http://s%d.b.hmny.io:9500"
+	TestNetWSPattern   = "ws://s%d.s.hmny.io:9800"
+)
+
 // PublicBlockChainAPI provides an API to access the Harmony blockchain.
 // It offers only methods that operate on public data that is freely available to anyone.
 type PublicBlockChainAPI struct {
@@ -64,61 +72,47 @@ func (s *PublicBlockChainAPI) GetBlockByHash(ctx context.Context, blockHash comm
 	return nil, err
 }
 
+// GenShardingStructure return sharding structure, given shard number and its patterns.
+func GenShardingStructure(shardNum, shardID uint32, httpPattern, wsPattern string) []map[string]interface{} {
+	res := []map[string]interface{}{}
+	for i := 0; i < int(shardNum); i++ {
+		res = append(res, map[string]interface{}{
+			"current": int(shardID) == i,
+			"shardID": i,
+			"http":    fmt.Sprintf(httpPattern, i),
+			"ws":      fmt.Sprintf(wsPattern, i),
+		})
+	}
+	return res
+}
+
 // GetShardingStructure returns an array of sharding structures.
 func (s *PublicBlockChainAPI) GetShardingStructure(ctx context.Context) ([]map[string]interface{}, error) {
+	// Get header and number of shards.
+	header := s.b.CurrentBlock().Header()
+	numShard := core.ShardingSchedule.InstanceForEpoch(header.Epoch).NumShards()
+	// Return shareding structure for each case.
 	if core.ShardingSchedule.GetNetworkID() == shardingconfig.MainNet {
-		return []map[string]interface{}{
+		s.b.CurrentBlock().Header()
-			map[string]interface{}{
+		return GenShardingStructure(numShard, s.b.GetShardID(), MainNetHTTPPattern, MainNetWSPattern), nil
-				"current": s.b.GetShardID() == 0,
-				"shardID": "0",
-				"http":    "http://s0.t.hmny.io:9500",
-				"ws":      "ws://s0.t.hmny.io:9800",
-			},
-			map[string]interface{}{
-				"current": s.b.GetShardID() == 1,
-				"shardID": "1",
-				"http":    "http://s1.t.hmny.io:9500",
-				"ws":      "ws://s1.t.hmny.io:9800",
-			},
-			map[string]interface{}{
-				"shardID": "2",
-				"http":    "http://s2.t.hmny.io:9500",
-				"ws":      "ws://s2.t.hmny.io:9800",
-			},
-			map[string]interface{}{
-				"shardID": "3",
-				"http":    "http://s3.t.hmny.io:9500",
-				"ws":      "ws://s3.t.hmny.io:9800",
-			},
-		}, nil
 	} else if core.ShardingSchedule.GetNetworkID() == shardingconfig.TestNet {
-		return []map[string]interface{}{
+		return GenShardingStructure(numShard, s.b.GetShardID(), TestNetHTTPPattern, TestNetWSPattern), nil
-			map[string]interface{}{
-				"shardID": "0",
-				"http":    "http://s0.b.hmny.io:9500",
-				"ws":      "ws://s0.s.hmny.io:9800",
-			},
-			map[string]interface{}{
-				"shardID": "1",
-				"http":    "http://s1.b.hmny.io:9500",
-				"ws":      "ws://s1.s.hmny.io:9800",
-			},
-		}, nil
 	} else {
 		return []map[string]interface{}{
 			map[string]interface{}{
-				"shardID": "0",
+				"current": s.b.GetShardID() == 0,
+				"shardID": 0,
 				"http":    "http://127.0.0.1:9500",
 				"ws":      "ws://127.0.0.1:9800",
 			},
 			map[string]interface{}{
-				"shardID": "1",
+				"current": s.b.GetShardID() == 1,
+				"shardID": 1,
 				"http":    "http://127.0.0.1:9501",
 				"ws":      "ws://127.0.0.1:9801",
 			},
 		}, nil
 	}
-
 }
 
 // GetCode returns the code stored at the given address in the state for the given block number.

internal/hmyapi/blockchain_test.go

@@ -0,0 +1,29 @@
+package hmyapi
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestGetShardingStructure(t *testing.T) {
+	shardID := 0
+	numShard := 4
+	res := GenShardingStructure(uint32(numShard), uint32(shardID), "http://s%d.t.hmy.io:9500", "ws://s%d.t.hmy.io:9800")
+	if len(res) != 4 || !res[0]["current"].(bool) || res[1]["current"].(bool) || res[2]["current"].(bool) || res[3]["current"].(bool) {
+		t.Error("Error when generating sharding structure")
+	}
+	for i := 0; i < numShard; i++ {
+		if res[i]["current"].(bool) != (i == shardID) {
+			t.Error("Error when generating sharding structure")
+		}
+		if res[i]["shardID"].(int) != i {
+			t.Error("Error when generating sharding structure")
+		}
+		if res[i]["http"].(string) != fmt.Sprintf("http://s%d.t.hmy.io:9500", i) {
+			t.Error("Error when generating sharding structure")
+		}
+		if res[i]["ws"].(string) != fmt.Sprintf("ws://s%d.t.hmy.io:9800", i) {
+			t.Error("Error when generating sharding structure")
+		}
+	}
+}