Fix tests

stakewise · Dec 29, 2024 · 21cb059 · 21cb059
1 parent 6260010
commit 21cb059
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 106 deletions.
diff --git a/src/validators/keystores/hashi_vault.py b/src/validators/keystores/hashi_vault.py
@@ -63,15 +63,18 @@ def secret_url(self, key_path: str, location: str = 'data') -> str:
 
 
 class HashiKeys:
+    """
+    Wrapper around Keys that proactively searches for duplicate keys to prevent
+    potential slashing.
+    """
+
     def __init__(self, keys: Keys | None = None):
         self.keys = keys or Keys({})
 
     def __getitem__(self, key: HexStr) -> BLSPrivkey:
         return self.keys[key]
 
     def __setitem__(self, key: HexStr, value: BLSPrivkey) -> None:
-        """Add new key/value pair proactively searching for duplicate keys to prevent
-        potential slashing."""
         if key in self.keys:
             raise RuntimeError(f'Duplicate validator key {key} found in hashi vault')
         self.keys[key] = value
@@ -83,6 +86,9 @@ def update(self, new_keys: 'HashiKeys' | Keys) -> None:
         for key, value in new_keys.items():
             self[key] = value
 
+    def __len__(self) -> int:
+        return len(self.keys)
+
     def __repr__(self) -> str:
         return f'HashiKeys({self.keys})'
 

diff --git a/src/validators/keystores/tests/test_hashi_vault.py b/src/validators/keystores/tests/test_hashi_vault.py
@@ -1,12 +1,9 @@
 import pytest
-from aiohttp.client import ClientSession
 
 from src.config.settings import settings
 from src.validators.keystores.hashi_vault import (
-    HashiVaultBundledKeysLoader,
     HashiVaultConfiguration,
     HashiVaultKeystore,
-    HashiVaultPrefixedKeysLoader,
 )
 
 
@@ -19,67 +16,14 @@ async def test_hashi_vault_bundled_keystores_loading(
         settings.hashi_vault_url = hashi_vault_url
         settings.hashi_vault_engine_name = 'secret'
         settings.hashi_vault_token = 'Secret'
-        settings.hashi_vault_key_paths = []
+        settings.hashi_vault_key_paths = ['ethereum/signing/keystores']
         settings.hashi_vault_key_prefixes = []
         settings.hashi_vault_parallelism = 1
 
-        config = HashiVaultConfiguration.from_settings()
-
-        async with ClientSession() as session:
-            keystore = await HashiVaultBundledKeysLoader._load_bundled_keys(
-                session=session,
-                secret_url=config.secret_url('ethereum/signing/keystores'),
-            )
+        keystore = await HashiVaultKeystore.load()
 
         assert len(keystore) == 2
 
-    @pytest.mark.usefixtures('mocked_hashi_vault')
-    async def test_hashi_vault_prefixed_keystores_finding(
-        self,
-        hashi_vault_url: str,
-    ):
-        settings.hashi_vault_url = hashi_vault_url
-        settings.hashi_vault_engine_name = 'secret'
-        settings.hashi_vault_token = 'Secret'
-        settings.hashi_vault_key_paths = []
-        settings.hashi_vault_key_prefixes = []
-        settings.hashi_vault_parallelism = 1
-
-        config = HashiVaultConfiguration.from_settings()
-
-        async with ClientSession() as session:
-            keystores_prefixes = await HashiVaultPrefixedKeysLoader._find_prefixed_hashi_vault_keys(
-                session=session,
-                prefix='ethereum/signing/prefixed1',
-                prefix_url=config.prefix_url('ethereum/signing/prefixed1'),
-            )
-        assert len(keystores_prefixes) == 2
-
-    @pytest.mark.usefixtures('mocked_hashi_vault')
-    async def test_hashi_vault_prefixed_keystores_loading(
-        self,
-        hashi_vault_url: str,
-    ):
-        settings.hashi_vault_url = hashi_vault_url
-        settings.hashi_vault_engine_name = 'secret'
-        settings.hashi_vault_token = 'Secret'
-        settings.hashi_vault_key_paths = []
-        settings.hashi_vault_key_prefixes = []
-        settings.hashi_vault_parallelism = 1
-
-        config = HashiVaultConfiguration.from_settings()
-
-        async with ClientSession() as session:
-            keystore = await HashiVaultPrefixedKeysLoader._load_prefixed_key(
-                session=session,
-                secret_url=config.secret_url(
-                    'ethereum/signing/prefixed1/8b09379ca969e8283a42a09285f430e8bd58c70bb33b44397ae81dac01b1403d0f631f156d211b6931a1c6284e2e469c',
-                ),
-            )
-        assert list(keystore.keys()) == [
-            '0x8b09379ca969e8283a42a09285f430e8bd58c70bb33b44397ae81dac01b1403d0f631f156d211b6931a1c6284e2e469c'
-        ]
-
     @pytest.mark.usefixtures('mocked_hashi_vault')
     async def test_hashi_vault_keystores_not_configured(
         self,
@@ -102,19 +46,14 @@ async def test_hashi_vault_bundled_keystores_inaccessible(
         settings.hashi_vault_url = hashi_vault_url
         settings.hashi_vault_engine_name = 'secret'
         settings.hashi_vault_token = 'Secret'
-        settings.hashi_vault_key_path = []
+        settings.hashi_vault_key_paths = ['ethereum/inaccessible/keystores']
         settings.hashi_vault_key_prefixes = []
         settings.hashi_vault_parallelism = 1
 
         with pytest.raises(
             RuntimeError, match='Can not retrieve validator signing keys from hashi vault'
         ):
-            config = HashiVaultConfiguration.from_settings()
-            async with ClientSession() as session:
-                await HashiVaultBundledKeysLoader._load_bundled_keys(
-                    session=session,
-                    secret_url=config.secret_url('ethereum/inaccessible/keystores'),
-                )
+            await HashiVaultKeystore.load()
 
     @pytest.mark.usefixtures('mocked_hashi_vault')
     async def test_hashi_vault_bundled_keystores_parallel(
@@ -131,15 +70,9 @@ async def test_hashi_vault_bundled_keystores_parallel(
         settings.hashi_vault_key_prefixes = []
         settings.hashi_vault_parallelism = 2
 
-        config = HashiVaultConfiguration.from_settings()
-        loader = HashiVaultBundledKeysLoader(
-            config=config,
-            input_iter=iter(settings.hashi_vault_key_paths),
-        )
-        keys = {}
-        await loader.load(keys)
+        keystore = await HashiVaultKeystore.load()
 
-        assert len(keys) == 4
+        assert len(keystore) == 4
 
     @pytest.mark.usefixtures('mocked_hashi_vault')
     async def test_hashi_vault_bundled_keystores_sequential(
@@ -155,16 +88,9 @@ async def test_hashi_vault_bundled_keystores_sequential(
         ]
         settings.hashi_vault_parallelism = 1
 
-        config = HashiVaultConfiguration.from_settings()
+        keystore = await HashiVaultKeystore.load()
 
-        loader = HashiVaultBundledKeysLoader(
-            config=config,
-            input_iter=iter(settings.hashi_vault_key_paths),
-        )
-        keys = {}
-        await loader.load(keys)
-
-        assert len(keys) == 4
+        assert len(keystore) == 4
 
     @pytest.mark.usefixtures('mocked_hashi_vault')
     async def test_hashi_vault_duplicates_parallel(
@@ -180,9 +106,8 @@ async def test_hashi_vault_duplicates_parallel(
         ]
         settings.hashi_vault_parallelism = 2
 
-        keystore = HashiVaultKeystore({})
-        with pytest.raises(RuntimeError, match='Found duplicate key in path'):
-            await keystore.load()
+        with pytest.raises(RuntimeError, match='Duplicate validator key'):
+            await HashiVaultKeystore.load()
 
     @pytest.mark.usefixtures('mocked_hashi_vault')
     async def test_hashi_vault_keystores_loading_custom_engine_name(
@@ -192,16 +117,10 @@ async def test_hashi_vault_keystores_loading_custom_engine_name(
         settings.hashi_vault_url = hashi_vault_url
         settings.hashi_vault_engine_name = 'custom'
         settings.hashi_vault_token = 'Secret'
-        settings.hashi_vault_key_paths = []
+        settings.hashi_vault_key_paths = ['ethereum/signing/keystores']
         settings.hashi_vault_parallelism = 1
 
-        config = HashiVaultConfiguration.from_settings()
-
-        async with ClientSession() as session:
-            keystore = await HashiVaultBundledKeysLoader._load_bundled_keys(
-                session=session,
-                secret_url=config.secret_url('ethereum/signing/keystores'),
-            )
+        keystore = await HashiVaultKeystore.load()
 
         assert len(keystore) == 2
 
@@ -214,16 +133,10 @@ async def test_hashi_vault_keystores_prefixed_loader(
         settings.hashi_vault_engine_name = 'secret'
         settings.hashi_vault_token = 'Secret'
         settings.hashi_vault_key_paths = []
-        settings.hashi_vault_key_prefixes = []
+        settings.hashi_vault_key_prefixes = ['ethereum/signing/prefixed1']
         settings.hashi_vault_parallelism = 1
 
-        config = HashiVaultConfiguration.from_settings()
-
-        loader = HashiVaultPrefixedKeysLoader(
-            config=config, input_iter=iter(['ethereum/signing/prefixed1'])
-        )
-        keystore = {}
-        await loader.load(keystore)
+        keystore = await HashiVaultKeystore.load()
 
         assert len(keystore) == 2
 
@@ -245,6 +158,6 @@ async def test_hashi_vault_load_bundled_and_prefixed(
         ]
         settings.hashi_vault_parallelism = 2
 
-        keystore = HashiVaultKeystore({})
-        keys = await keystore.load()
-        assert len(keys) == 8
+        keystore = await HashiVaultKeystore.load()
+
+        assert len(keystore) == 8