[py]: repr for script keys, types and nicer error when unpinning unknown keys and unit tests

symonk · symonk · commit 4440ea0bd6a2 · 2022-08-14T10:02:13.000+01:00
diff --git a/py/selenium/webdriver/remote/script_key.py b/py/selenium/webdriver/remote/script_key.py
@@ -28,3 +28,6 @@ def id(self):
 
     def __eq__(self, other):
         return self._id == other
+
+    def __repr__(self) -> str:
+        return f"ScriptKey(id={self.id})"
diff --git a/py/selenium/webdriver/remote/webdriver.py b/py/selenium/webdriver/remote/webdriver.py
@@ -18,48 +18,42 @@
 """The WebDriver implementation."""
 import contextlib
 import copy
+import pkgutil
 import types
 import typing
-from importlib import import_module
-
-import pkgutil
-
-from typing import Dict, List, Optional, Union
-
 import warnings
-
 from abc import ABCMeta
 from base64 import b64decode, urlsafe_b64encode
 from contextlib import asynccontextmanager, contextmanager
-
-from .bidi_connection import BidiConnection
-from .command import Command
-from .errorhandler import ErrorHandler
-from .file_detector import FileDetector, LocalFileDetector
-from .mobile import Mobile
-from .remote_connection import RemoteConnection
-from .script_key import ScriptKey
-from .shadowroot import ShadowRoot
-from .switch_to import SwitchTo
-from .webelement import WebElement
+from importlib import import_module
+from typing import Dict, List, Optional, Union
 
 from selenium.common.exceptions import (InvalidArgumentException,
                                         JavascriptException,
                                         WebDriverException,
                                         NoSuchCookieException,
                                         NoSuchElementException)
 from selenium.webdriver.common.by import By
+from selenium.webdriver.common.html5.application_cache import ApplicationCache
 from selenium.webdriver.common.options import BaseOptions
 from selenium.webdriver.common.print_page_options import PrintOptions
 from selenium.webdriver.common.timeouts import Timeouts
-from selenium.webdriver.common.html5.application_cache import ApplicationCache
-from selenium.webdriver.support.relative_locator import RelativeBy
 from selenium.webdriver.common.virtual_authenticator import (
     Credential,
     VirtualAuthenticatorOptions,
     required_virtual_authenticator
 )
-
+from selenium.webdriver.support.relative_locator import RelativeBy
+from .bidi_connection import BidiConnection
+from .command import Command
+from .errorhandler import ErrorHandler
+from .file_detector import FileDetector, LocalFileDetector
+from .mobile import Mobile
+from .remote_connection import RemoteConnection
+from .script_key import ScriptKey
+from .shadowroot import ShadowRoot
+from .switch_to import SwitchTo
+from .webelement import WebElement
 
 _W3C_CAPABILITY_NAMES = frozenset([
     'acceptInsecureCerts',
@@ -456,16 +450,21 @@ def title(self) -> str:
         """
         return self.execute(Command.GET_TITLE).get("value", "")
 
-    def pin_script(self, script, script_key=None) -> ScriptKey:
-        _script_key = ScriptKey(script_key)
-        self.pinned_scripts[_script_key.id] = script
-        return _script_key
+    def pin_script(self, script: str, script_key=None) -> ScriptKey:
+        """Store common javascript scripts to be executed later by a unique hashable ID."""
+        script_key_instance = ScriptKey(script_key)
+        self.pinned_scripts[script_key_instance.id] = script
+        return script_key_instance
 
-    def unpin(self, script_key) -> None:
-        self.pinned_scripts.pop(script_key.id)
+    def unpin(self, script_key: ScriptKey) -> None:
+        """Remove a pinned script from storage."""
+        try:
+            self.pinned_scripts.pop(script_key.id)
+        except KeyError:
+            raise KeyError(f"No script with key: {script_key} existed in {self.pinned_scripts}") from None
 
     def get_pinned_scripts(self) -> List[str]:
-        return list(self.pinned_scripts.keys())
+        return list(self.pinned_scripts)
 
     def execute_script(self, script, *args):
         """
