chore: suggestions are implemented

ApeWorX · Aug 7, 2024 · aa62387 · aa62387
1 parent 8cb8709
commit aa62387
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 30 deletions.
diff --git a/ape_infura/provider.py b/ape_infura/provider.py
@@ -44,14 +44,13 @@ def __init__(self, *args, **kwargs):
     def load_api_keys(self):
         self.api_keys = []
         for env_var_name in _ENVIRONMENT_VARIABLE_NAMES:
-            env_var = os.environ.get(env_var_name)
-            if env_var:
+            if env_var := os.environ.get(env_var_name):
                 self.api_keys.extend([key.strip() for key in env_var.split(",")])
 
         if not self.api_keys:
             raise MissingProjectKeyError()
 
-    def get_random_api_key(self):
+    def get_random_api_key(self) -> str:
         return random.choice(self.api_keys)
 
     @property

diff --git a/tests/test_provider.py b/tests/test_provider.py
@@ -1,5 +1,4 @@
 import os
-from unittest.mock import patch
 
 import pytest
 import websocket  # type: ignore
@@ -36,42 +35,42 @@ def test_infura_ws(provider):
         pytest.fail(f"Websocket URI not accessible. Reason: {err}")
 
 
-def test_load_multiple_api_keys(provider):
-    with patch.dict(
+def test_load_multiple_api_keys(provider, mocker):
+    mocker.patch.dict(
         os.environ,
         {"WEB3_INFURA_PROJECT_ID": "key1,key2,key3", "WEB3_INFURA_API_KEY": "key4,key5,key6"},
-    ):
-        provider.load_api_keys()
-        # As there will be API keys in the ENV as well
-        assert len(provider.api_keys) >= 6
-        assert "key1" in provider.api_keys
-        assert "key6" in provider.api_keys
+    )
+    provider.load_api_keys()
+    # As there will be API keys in the ENV as well
+    assert len(provider.api_keys) == 6
+    assert "key1" in provider.api_keys
+    assert "key6" in provider.api_keys
 
 
-def test_load_single_and_multiple_api_keys(provider):
-    with patch.dict(
+def test_load_single_and_multiple_api_keys(provider, mocker):
+    mocker.patch.dict(
         os.environ,
         {
             "WEB3_INFURA_PROJECT_ID": "single_key1",
             "WEB3_INFURA_API_KEY": "single_key2",
         },
-    ):
-        provider.load_api_keys()
-        assert len(provider.api_keys) >= 2
-        assert "single_key1" in provider.api_keys
-        assert "single_key2" in provider.api_keys
+    )
+    provider.load_api_keys()
+    assert len(provider.api_keys) == 2
+    assert "single_key1" in provider.api_keys
+    assert "single_key2" in provider.api_keys
 
 
-def test_random_api_key_selection(provider):
-    with patch.dict(os.environ, {"WEB3_INFURA_PROJECT_ID": "key1,key2,key3,key4,key5"}):
-        provider.load_api_keys()
-        selected_keys = set()
-        for _ in range(50):  # Run multiple times to ensure randomness
-            selected_keys.add(provider.get_random_api_key())
-        assert len(selected_keys) > 1  # Ensure we're getting different keys
+def test_random_api_key_selection(provider, mocker):
+    mocker.patch.dict(os.environ, {"WEB3_INFURA_PROJECT_ID": "key1,key2,key3,key4,key5"})
+    provider.load_api_keys()
+    selected_keys = set()
+    for _ in range(50):  # Run multiple times to ensure randomness
+        selected_keys.add(provider.get_random_api_key())
+    assert len(selected_keys) > 1  # Ensure we're getting different keys
 
 
-def test_missing_project_key_error_raised(provider):
-    with patch.dict(os.environ, {}, clear=True):
-        with pytest.raises(MissingProjectKeyError):
-            provider.load_api_keys()
+def test_missing_project_key_error_raised(provider, mocker):
+    mocker.patch.dict(os.environ, {}, clear=True)
+    with pytest.raises(MissingProjectKeyError):
+        provider.load_api_keys()