From 72d0a5a538cd74b0adc7ba178abf5a08b887a50c Mon Sep 17 00:00:00 2001
From: wacban <wacban@users.noreply.github.com>
Date: Tue, 12 Dec 2023 11:57:59 +0000
Subject: [PATCH] test(resharding) - test for getting the shard layout info by
 rpc (#10320)

We need a way for near ecosystem developers to query the current shard
layout. The EXPERIMENTAL_protocol_config allows that although with some
extra steps. It should be good enough until we have time to implement a
dedicated endpoint just for the shard layout. In this PR I extended the
resharding nayduck test to check that the mentioned rpc endpoint works
as expected and that changes in the shard layout are reflected in the
rpc.
---
 pytest/lib/utils.py               |  2 --
 pytest/tests/sanity/resharding.py | 29 +++++++++++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/pytest/lib/utils.py b/pytest/lib/utils.py
index 42b0113e0ba..11f34855d00 100644
--- a/pytest/lib/utils.py
+++ b/pytest/lib/utils.py
@@ -1,4 +1,3 @@
-import atexit
 import base58
 import hashlib
 import json
@@ -7,7 +6,6 @@
 import random
 import re
 import shutil
-import subprocess
 import sys
 import tempfile
 import time
diff --git a/pytest/tests/sanity/resharding.py b/pytest/tests/sanity/resharding.py
index 0eaf5cb3b73..a6997a9fe94 100644
--- a/pytest/tests/sanity/resharding.py
+++ b/pytest/tests/sanity/resharding.py
@@ -71,15 +71,40 @@ def test_resharding(self):
 
         metrics_tracker = MetricsTracker(node0)
 
-        for height, _ in poll_blocks(node0):
+        for height, hash in poll_blocks(node0):
             version = metrics_tracker.get_int_metric_value(
                 "near_shard_layout_version")
             num_shards = metrics_tracker.get_int_metric_value(
                 "near_shard_layout_num_shards")
 
+            protocol_config = node0.json_rpc(
+                "EXPERIMENTAL_protocol_config",
+                {"block_id": hash},
+            )
+
+            self.assertTrue('error' not in protocol_config)
+
+            self.assertTrue('result' in protocol_config)
+            protocol_config = protocol_config.get('result')
+
+            self.assertTrue('shard_layout' in protocol_config)
+            shard_layout = protocol_config.get('shard_layout')
+
+            self.assertTrue('V1' in shard_layout)
+            shard_layout = shard_layout.get('V1')
+
+            self.assertTrue('boundary_accounts' in shard_layout)
+            boundary_accounts = shard_layout.get('boundary_accounts')
+
             logger.info(
-                f"#{height} shard layout version: {version}, num shards: {num_shards} "
+                f"#{height} shard layout version: {version}, num shards: {num_shards}"
             )
+            logger.debug(f"#{height} shard layout: {shard_layout}")
+
+            # check the shard layout versions from metrics and from json rpc are equal
+            self.assertEqual(version, shard_layout.get('version'))
+            # check the shard num from metrics and json rpc are equal
+            self.assertEqual(num_shards, len(boundary_accounts) + 1)
 
             # This may be flaky - it shouldn't - but it may. We collect metrics
             # after the block is processed. If there is some delay the shard