From 18074581c741f7b3efd21cc7524d0dd2375c5f1a Mon Sep 17 00:00:00 2001
From: Alex <18387287+wadealexc@users.noreply.github.com>
Date: Wed, 31 Jan 2024 14:07:32 -0500
Subject: [PATCH] test: fix flaky tests by removing bogosort (#163)

* test: fix flaky tests by removing bogosort

* test: fix flaky test by rejecting empty addr inputs
---
 test/unit/BLSApkRegistryUnit.t.sol      |  3 ++-
 test/unit/BLSSignatureCheckerUnit.t.sol |  2 +-
 test/utils/BLSMockAVSDeployer.sol       | 25 +++++++++++++++++++------
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/test/unit/BLSApkRegistryUnit.t.sol b/test/unit/BLSApkRegistryUnit.t.sol
index c85927be..82ff0a87 100644
--- a/test/unit/BLSApkRegistryUnit.t.sol
+++ b/test/unit/BLSApkRegistryUnit.t.sol
@@ -327,13 +327,14 @@ contract BLSApkRegistryUnitTests_registerBLSPublicKey is BLSApkRegistryUnitTests
         address operator,
         address operator2
     ) public {
+        cheats.assume(operator != address(0));
         cheats.assume(operator != operator2);
         BN254.G1Point memory messageHash =
             registryCoordinator.pubkeyRegistrationMessageHash(operator);
         pubkeyRegistrationParams.pubkeyRegistrationSignature = _signMessage(operator);
 
         cheats.startPrank(address(registryCoordinator));
-        blsApkRegistry.registerBLSPublicKey(operator, pubkeyRegistrationParams, messageHash);
+        blsApkRegistry.registerBLSPublicKey(operator, pubkeyRegistrationParams, messageHash);        
 
         cheats.expectRevert("BLSApkRegistry.registerBLSPublicKey: public key already registered");
         blsApkRegistry.registerBLSPublicKey(operator2, pubkeyRegistrationParams, messageHash);
diff --git a/test/unit/BLSSignatureCheckerUnit.t.sol b/test/unit/BLSSignatureCheckerUnit.t.sol
index f8479ccc..58c00c4b 100644
--- a/test/unit/BLSSignatureCheckerUnit.t.sol
+++ b/test/unit/BLSSignatureCheckerUnit.t.sol
@@ -38,7 +38,7 @@ contract BLSSignatureCheckerUnitTests is BLSMockAVSDeployer {
     // this test checks that a valid signature from maxOperatorsToRegister with a random number of nonsigners is checked
     // correctly on the BLSSignatureChecker contract when all operators are only regsitered for a single quorum and
     // the signature is only checked for stakes on that quorum
-    function test_checkSignatures_SingleQuorum(uint256 pseudoRandomNumber) public { 
+    function testFuzz_checkSignatures_SingleQuorum(uint256 pseudoRandomNumber) public { 
         uint256 numNonSigners = pseudoRandomNumber % (maxOperatorsToRegister - 1);
         uint256 quorumBitmap = 1;
         bytes memory quorumNumbers = BitmapUtils.bitmapToBytesArray(quorumBitmap);
diff --git a/test/utils/BLSMockAVSDeployer.sol b/test/utils/BLSMockAVSDeployer.sol
index 4e1476b6..7e870a7a 100644
--- a/test/utils/BLSMockAVSDeployer.sol
+++ b/test/utils/BLSMockAVSDeployer.sol
@@ -57,13 +57,22 @@ contract BLSMockAVSDeployer is MockAVSDeployer {
         uint256[] memory nonSignerPrivateKeys = new uint256[](numNonSigners);
         for (uint i = 0; i < numNonSigners; i++) {
             nonSignerPrivateKeys[i] = uint256(keccak256(abi.encodePacked("nonSignerPrivateKey", pseudoRandomNumber, i))) % BN254.FR_MODULUS;
-            uint256 j = 0;
-            if(i != 0) {
-                while(BN254.generatorG1().scalar_mul(nonSignerPrivateKeys[i]).hashG1Point() <= BN254.generatorG1().scalar_mul(nonSignerPrivateKeys[i-1]).hashG1Point()){
-                    nonSignerPrivateKeys[i] = uint256(keccak256(abi.encodePacked("nonSignerPrivateKey", pseudoRandomNumber, j))) % BN254.FR_MODULUS;
-                    j++;
-                }
+        }
+
+        // Sort nonSignerPrivateKeys in order of ascending pubkeyHash
+        // Uses insertion sort to sort array in place
+        for (uint i = 1; i < nonSignerPrivateKeys.length; i++) {
+            uint privateKey = nonSignerPrivateKeys[i];
+            bytes32 pubkeyHash = _toPubkeyHash(privateKey);
+            uint j = i;
+
+            // Move elements of nonSignerPrivateKeys[0..i-1] that are greater than the current key
+            // to one position ahead of their current position
+            while (j > 0 && _toPubkeyHash(nonSignerPrivateKeys[j - 1]) > pubkeyHash) {
+                nonSignerPrivateKeys[j] = nonSignerPrivateKeys[j - 1];
+                j--;
             }
+            nonSignerPrivateKeys[j] = privateKey;
         }
 
         return (signerPrivateKeys, nonSignerPrivateKeys);
@@ -135,4 +144,8 @@ contract BLSMockAVSDeployer is MockAVSDeployer {
 
         return (referenceBlockNumber, nonSignerStakesAndSignature);
     }
+
+    function _toPubkeyHash(uint privKey) internal view returns (bytes32) {
+        return BN254.generatorG1().scalar_mul(privKey).hashG1Point();
+    }
 }