From bc55a3a0c752ead8ca24a55fa23864f1988944ec Mon Sep 17 00:00:00 2001
From: Agnish Ghosh <80243668+agnxsh@users.noreply.github.com>
Date: Sun, 12 Jan 2025 10:34:55 +0530
Subject: [PATCH] add: `compute_columns_for_custody_group` tests (#6838)

* add new fixture networking tests

* update all tests

* fix test

* avoid extra variables

* other fixes

* oops

* updated tests file
---
 AllTests-mainnet.md                           |  9 ++++-
 beacon_chain/spec/peerdas_helpers.nim         |  4 +-
 .../test_fixture_networking.nim               | 40 +++++++++++++++----
 3 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/AllTests-mainnet.md b/AllTests-mainnet.md
index 76753732d..592997cf8 100644
--- a/AllTests-mainnet.md
+++ b/AllTests-mainnet.md
@@ -513,6 +513,11 @@ OK: 253/253 Fail: 0/253 Skip: 0/253
 OK: 56/56 Fail: 0/56 Skip: 0/56
 ## EF - PeerDAS - Networking [Preset: mainnet]
 ```diff
++ Networking - Compute Columns for Custody Group - mainnet/fulu/networking/compute_columns_f OK
++ Networking - Compute Columns for Custody Group - mainnet/fulu/networking/compute_columns_f OK
++ Networking - Compute Columns for Custody Group - mainnet/fulu/networking/compute_columns_f OK
++ Networking - Compute Columns for Custody Group - mainnet/fulu/networking/compute_columns_f OK
++ Networking - Compute Columns for Custody Group - mainnet/fulu/networking/compute_columns_f OK
 + Networking - Get Custody Groups - mainnet/fulu/networking/get_custody_groups/pyspec_tests/ OK
 + Networking - Get Custody Groups - mainnet/fulu/networking/get_custody_groups/pyspec_tests/ OK
 + Networking - Get Custody Groups - mainnet/fulu/networking/get_custody_groups/pyspec_tests/ OK
@@ -523,7 +528,7 @@ OK: 56/56 Fail: 0/56 Skip: 0/56
 + Networking - Get Custody Groups - mainnet/fulu/networking/get_custody_groups/pyspec_tests/ OK
 + Networking - Get Custody Groups - mainnet/fulu/networking/get_custody_groups/pyspec_tests/ OK
 ```
-OK: 9/9 Fail: 0/9 Skip: 0/9
+OK: 14/14 Fail: 0/14 Skip: 0/14
 ## EF - SSZ generic types
 ```diff
   Testing basic_vector inputs - invalid                                                      Skip
@@ -1149,4 +1154,4 @@ OK: 2/2 Fail: 0/2 Skip: 0/2
 OK: 9/9 Fail: 0/9 Skip: 0/9
 
 ---TOTAL---
-OK: 778/783 Fail: 0/783 Skip: 5/783
+OK: 783/788 Fail: 0/788 Skip: 5/788
diff --git a/beacon_chain/spec/peerdas_helpers.nim b/beacon_chain/spec/peerdas_helpers.nim
index 6faa3b62c..0f5acb2cd 100644
--- a/beacon_chain/spec/peerdas_helpers.nim
+++ b/beacon_chain/spec/peerdas_helpers.nim
@@ -25,8 +25,8 @@ type
   ProofBytes = array[fulu.CELLS_PER_EXT_BLOB, KzgProof]
 
 # https://github.com/ethereum/consensus-specs/blob/v1.5.0-alpha.10/specs/fulu/das-core.md#compute_columns_for_custody_group
-iterator compute_columns_for_custody_group(custody_group: CustodyIndex):
-                                           ColumnIndex =
+iterator compute_columns_for_custody_group*(custody_group: CustodyIndex):
+                                            ColumnIndex =
   for i in 0'u64 ..< COLUMNS_PER_GROUP:
     yield ColumnIndex(NUMBER_OF_CUSTODY_GROUPS * i + custody_group)
 
diff --git a/tests/consensus_spec/test_fixture_networking.nim b/tests/consensus_spec/test_fixture_networking.nim
index f02d573ec..95dbadd01 100644
--- a/tests/consensus_spec/test_fixture_networking.nim
+++ b/tests/consensus_spec/test_fixture_networking.nim
@@ -18,9 +18,28 @@ import
   ../testutil,
   ./fixtures_utils, ./os_ops
 
-from std/sequtils import mapIt
+proc runComputeForCustodyGroup(suiteName, path: string) =
+  let relativeTestPathComponent = path.relativeTestPathComponent()
+  test "Networking - Compute Columns for Custody Group - " &
+      relativeTestPathComponent:
+    type TestMetaYaml = object
+      custody_group: uint64
+      result: seq[uint64]
+    let
+      meta = block:
+        var s = openFileStream(path/"meta.yaml")
+        defer: close(s)
+        var res: TestMetaYaml
+        yaml.load(s, res)
+        res
+      custody_group = meta.custody_group
 
-proc runGetCustodyColumns(suiteName, path: string) =
+    var counter = 0
+    for column in compute_columns_for_custody_group(custody_group):
+      check column == meta.result[counter]
+      inc counter
+
+proc runGetCustodyGroups(suiteName, path: string) =
   let relativePathComponent = path.relativeTestPathComponent()
   test "Networking - Get Custody Groups - " & relativePathComponent:
     type TestMetaYaml = object
@@ -36,17 +55,22 @@ proc runGetCustodyColumns(suiteName, path: string) =
         res
       node_id = UInt256.fromDecimal(meta.node_id)
       custody_group_count = meta.custody_group_count
-      reslt = (meta.result).mapIt(it)
 
     let columns = get_custody_groups(node_id, custody_group_count)
 
     for i in 0..<columns.lenu64:
-      check columns[i] == reslt[i]
+      check columns[i] == meta.result[i]
 
 suite "EF - PeerDAS - Networking" & preset():
   const presetPath = SszTestsDir/const_preset
   # foldering to be resolved in alpha 11 release of consensus spec tests
-  let basePath =
-    presetPath/"fulu"/"networking"/"get_custody_groups"/"pyspec_tests"
-  for kind, path in walkDir(basePath, relative = true, checkDir = true):
-    runGetCustodyColumns(suiteName, basePath/path)
\ No newline at end of file
+  block:
+    let basePath =
+      presetPath/"fulu"/"networking"/"get_custody_groups"/"pyspec_tests"
+    for kind, path in walkDir(basePath, relative = true, checkDir = true):
+      runGetCustodyGroups(suiteName, basePath/path)
+  block:
+    let basePath =
+      presetPath/"fulu"/"networking"/"compute_columns_for_custody_group"/"pyspec_tests"
+    for kind, path in walkDir(basePath, relative = true, checkDir = true):
+      runComputeForCustodyGroup(suiteName, basePath/path)