Tests / code cleanup

src/epplibwrapper/client.py

@@ -44,7 +44,7 @@ class EPPLibWrapper:
     ATTN: This should not be used directly. Use `Domain` from domain.py.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, start_connection_pool=True) -> None:
         """Initialize settings which will be used for all connections."""
         # prepare (but do not send) a Login command
         self._login = commands.Login(
@@ -80,7 +80,8 @@ class EPPLibWrapper:
         # Tracks the status of the pool
         self.pool_status = PoolStatus()
 
-        self.start_connection_pool()
+        if start_connection_pool:
+            self.start_connection_pool()
 
     def _send(self, command):
         """Helper function used by `send`."""
@@ -103,21 +104,21 @@ class EPPLibWrapper:
                 self.start_connection_pool()
         except (ValueError, ParsingError) as err:
             message = f"{cmd_type} failed to execute due to some syntax error."
-            logger.warning(message, exc_info=True)
+            logger.error(message, exc_info=True)
             raise RegistryError(message) from err
         except TransportError as err:
             message = f"{cmd_type} failed to execute due to a connection error."
-            logger.warning(message, exc_info=True)
+            logger.error(message, exc_info=True)
             raise RegistryError(message) from err
         except LoginError as err:
             # For linter
             text = "failed to execute due to a registry login error."
             message = f"{cmd_type} {text}"
-            logger.warning(message, exc_info=True)
+            logger.error(message, exc_info=True)
             raise RegistryError(message) from err
         except Exception as err:
             message = f"{cmd_type} failed to execute due to an unknown error."
-            logger.warning(message, exc_info=True)
+            logger.error(message, exc_info=True)
             raise RegistryError(message) from err
         else:
             if response.code >= 2000:
@@ -134,15 +135,15 @@ class EPPLibWrapper:
             raise ValueError("Please sanitize user input before sending it.")
 
         # Reopen the pool if its closed
+        # Only occurs when a login error is raised, after connection is successful
         if not self.pool_status.pool_running:
             # We want to reopen the connection pool,
             # but we don't want the end user to wait while it opens.
             # Raise syntax doesn't allow this, so we use a try/catch
             # block.
             try:
-                raise RegistryError(
+                logger.error("Can't contact the Registry. Pool was not running.")
-                    "Can't contact the Registry. Please try again later"
+                raise RegistryError("Can't contact the Registry. Pool was not running.")
-                )
             except RegistryError as err:
                 raise err
             finally:
@@ -159,39 +160,46 @@ class EPPLibWrapper:
                 else:  # don't try again
                     raise err
 
-    def start_connection_pool(self, restart_pool_if_exists=True):
+    def start_connection_pool(
+        self, restart_pool_if_exists=True, try_start_if_invalid=False
+    ):
         """Starts a connection pool for the registry.
 
         restart_pool_if_exists -> bool:
         If an instance of the pool already exists,
         then then that instance will be killed first.
-        It is generally recommended to keep this enabled."""
+        It is generally recommended to keep this enabled.
+
+        try_start_if_invalid -> bool:
+        Designed for use in test cases, if we can't connect
+        to the registry, ignore that and try to connect anyway
+        It is generally recommended to keep this disabled.
+        """
         # Since we reuse the same creds for each pool, we can test on
         # one socket, and if successful, then we know we can connect.
-        if settings.DEBUG or self._test_registry_connection_success():
+        if (
+            not try_start_if_invalid
+            and settings.DEBUG
+            or not self._test_registry_connection_success()
+        ):
             logger.warning("Cannot contact the Registry")
             self.pool_status.connection_success = False
-            # Q: Should err be raised instead?
-            # Q2: Since we try to connect 3 times,
-            # this indicates that the Registry isn't responsive.
-            # What should we do in this case?
-            return
         else:
             self.pool_status.connection_success = True
 
-        # If this function is reinvoked, then ensure
+            # If this function is reinvoked, then ensure
-        # that we don't have duplicate data sitting around.
+            # that we don't have duplicate data sitting around.
-        if self._pool is not None and restart_pool_if_exists:
+            if self._pool is not None and restart_pool_if_exists:
-            logger.info("Connection pool restarting...")
+                logger.info("Connection pool restarting...")
-            self.kill_pool()
+                self.kill_pool()
 
-        self._pool = EPPConnectionPool(
+            self._pool = EPPConnectionPool(
-            client=self._client, login=self._login, options=self.pool_options
+                client=self._client, login=self._login, options=self.pool_options
-        )
+            )
-        self.pool_status.pool_running = True
+            self.pool_status.pool_running = True
-        self.pool_status.pool_hanging = False
+            self.pool_status.pool_hanging = False
 
-        logger.info("Connection pool started")
+            logger.info("Connection pool started")
 
     def kill_pool(self):
         """Kills the existing pool. Use this instead
@@ -220,7 +228,7 @@ class EPPLibWrapper:
 try:
     # Initialize epplib
     CLIENT = EPPLibWrapper()
-    logger.debug("registry client initialized")
+    logger.info("registry client initialized")
 except Exception:
     CLIENT = None  # type: ignore
     logger.warning(

src/epplibwrapper/socket.py

@@ -16,7 +16,7 @@ logger = logging.getLogger(__name__)
 class Socket:
     """Context manager which establishes a TCP connection with registry."""
 
-    def __init__(self, client: commands.Login, login: Client) -> None:
+    def __init__(self, client: Client, login: commands.Login) -> None:
         """Save the epplib client and login details."""
         self.client = client
         self.login = login
@@ -29,7 +29,7 @@ class Socket:
         """Runs disconnect(), which closes a connection with EPPLib."""
         self.disconnect()
 
-    def connect(self, pass_response_only=False):
+    def connect(self):
         """Use epplib to connect."""
         self.client.connect()
         response = self.client.send(self.login)

src/epplibwrapper/tests/test_pool.py

@@ -1,10 +1,14 @@
 from unittest import skip
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 from django.conf import settings
 
 from django.test import Client
+from django.test import TestCase
 from epplibwrapper.client import EPPLibWrapper
+from epplibwrapper.utility.pool import EPPConnectionPool
+from registrar.models.domain import Domain
 from registrar.tests.common import MockEppLib
+from registrar.models.domain import registry
 
 import logging
 
@@ -12,24 +16,17 @@ try:
     from epplib.client import Client
     from epplib import commands
     from epplib.exceptions import TransportError
-    from epplib.transport import SocketTransport
+    from epplib.responses import base
 except ImportError:
     pass
 
 logger = logging.getLogger(__name__)
 
 
-@patch("djangooidc.views.CLIENT", autospec=True)
+class TestConnectionPool(TestCase):
-class TestConnectionPool(MockEppLib):
     """Tests for our connection pooling behaviour"""
 
     def setUp(self):
-        """
-        Background:
-            Given the registrant is logged in
-            And the registrant is the admin on a domain
-        """
-        super().setUp()
         self.pool_options = {
             # Current pool size
             "size": 1,
@@ -40,10 +37,45 @@ class TestConnectionPool(MockEppLib):
             "keepalive": 60,
         }
 
-    def tearDown(self):
+        # Mock a successful connection
-        super().tearDown()
+        self.mock_connect_patch = patch("epplib.client.Client.connect")
+        self.mocked_connect_function = self.mock_connect_patch.start()
+        self.mocked_connect_function.side_effect = self.mock_connect
 
-    def user_info(*args):
+        # Mock the send behaviour
+        self.mock_send_patch = patch("epplib.client.Client.send")
+        self.mocked_send_function = self.mock_send_patch.start()
+        self.mocked_send_function.side_effect = self.mock_send
+
+        # Mock the pool object
+        self.mockSendPatch = patch("registrar.models.domain.registry._pool")
+        self.mockedSendFunction = self.mockSendPatch.start()
+        self.mockedSendFunction.side_effect = self.fake_pool
+
+    def tearDown(self):
+        self.mock_send_patch.stop()
+        self.mock_connect_patch.stop()
+        self.mockSendPatch.stop()
+
+    def mock_connect(self, _request):
+        return None
+
+    def mock_send(self, _request):
+        if isinstance(_request, commands.Login):
+            response = MagicMock(
+                code=1000,
+                msg="Command completed successfully",
+                res_data=None,
+                cl_tr_id="xkw1uo#2023-10-17T15:29:09.559376",
+                sv_tr_id="5CcH4gxISuGkq8eqvr1UyQ==-35a",
+                extensions=[],
+                msg_q=None,
+            )
+
+            return response
+        return None
+
+    def user_info(self, *args):
         return {
             "sub": "TEST",
             "email": "test@example.com",
@@ -52,67 +84,19 @@ class TestConnectionPool(MockEppLib):
             "phone": "814564000",
         }
 
-    def test_pool_created_successfully(self, mock_client):
+    @patch("djangooidc.views.CLIENT", autospec=True)
-        # setup
+    def fake_pool(self, mock_client):
-        session = self.client.session
+        # mock client
-        session["state"] = "TEST"  # nosec B105
-        session.save()
-        # mock
         mock_client.callback.side_effect = self.user_info
+        # Create a mock transport object
+        mock_login = MagicMock()
+        mock_login.cert_file = "path/to/cert_file"
+        mock_login.key_file = "path/to/key_file"
 
-        client = EPPLibWrapper()
+        pool = EPPConnectionPool(
-        pool = client._pool
+            client=mock_client, login=mock_login, options=self.pool_options
-
-        # These are defined outside of the pool,
-        # so we can reimplement how this is being done
-        # in client.py. They should remain unchanged,
-        # and if they aren't, something went wrong.
-        expected_login = commands.Login(
-            cl_id="nothing",
-            password="nothing",
-            obj_uris=[
-                "urn:ietf:params:xml:ns:domain-1.0",
-                "urn:ietf:params:xml:ns:contact-1.0",
-            ],
-            new_pw=None,
-            version="1.0",
-            lang="en",
-            ext_uris=[],
         )
-
+        return pool
-        # Key/cert will generate a new file everytime.
-        # This should never be null, so we can check for that.
-        try:
-            expected_client = Client(
-                SocketTransport(
-                    settings.SECRET_REGISTRY_HOSTNAME,
-                    cert_file=pool._client.transport.cert_file,
-                    key_file=pool._client.transport.key_file,
-                    password=settings.SECRET_REGISTRY_KEY_PASSPHRASE,
-                )
-            ).__dict__
-        except Exception as err:
-            self.fail(err)
-
-        # We don't care about checking if the objects are both of
-        # the same reference, we only care about data parity, so
-        # we do a dict conversion.
-        actual_client = pool._client.__dict__
-        actual_client["transport"] = actual_client["transport"].__dict__
-        expected_client["transport"] = expected_client["transport"].__dict__
-
-        # Ensure that we're getting the credentials we expect
-        self.assertEqual(pool._login, expected_login)
-        self.assertEqual(actual_client, expected_client)
-
-        # Check that options are set correctly
-        self.assertEqual(pool.size, self.pool_options["size"])
-        self.assertEqual(pool.keepalive, self.pool_options["keepalive"])
-        self.assertEqual(pool.exc_classes, self.pool_options["exc_classes"])
-
-        # Check that it is running
-        self.assertEqual(client.pool_status.connection_success, True)
-        self.assertEqual(client.pool_status.pool_running, True)
 
     @skip("not implemented yet")
     def test_pool_timesout(self):
@@ -124,7 +108,33 @@ class TestConnectionPool(MockEppLib):
         """Multiple users send data concurrently"""
         raise
 
-    @skip("not implemented yet")
+    def test_pool_tries_create_invalid(self):
+        """A .send is invoked on the pool, but the pool
+        shouldn't be running."""
+        # Fake data for the _pool object
+        domain, _ = Domain.objects.get_or_create(name="freeman.gov")
+
+        # Trigger the getter - should fail
+        expected_contact = domain.security_contact
+        self.assertEqual(registry.pool_status.pool_running, False)
+        self.assertEqual(registry.pool_status.connection_success, False)
+        self.assertEqual(len(registry._pool.conn), 0)
+
     def test_pool_sends_data(self):
-        """A .send is invoked on the pool"""
+        """A .send is invoked on the pool successfully"""
-        raise
+        # Fake data for the _pool object
+        domain, _ = Domain.objects.get_or_create(name="freeman.gov")
+
+        # The connection pool will fail to start, start it manually
+        # so that our mocks can take over
+        registry.start_connection_pool(try_start_if_invalid=True)
+
+        # Pretend that we've connected
+        registry.pool_status.pool_running = True
+        registry.pool_status.connection_success = True
+
+        # Trigger the getter - should succeed
+        expected_contact = domain.security_contact
+        self.assertEqual(registry.pool_status.pool_running, True)
+        self.assertEqual(registry.pool_status.connection_success, True)
+        self.assertEqual(len(registry._pool.conn), self.pool_options["size"])

src/epplibwrapper/utility/pool.py

@@ -1,8 +1,8 @@
 import logging
 import gevent
 from geventconnpool import ConnectionPool
-from epplibwrapper.errors import RegistryError, LoginError
 from epplibwrapper.socket import Socket
+from epplibwrapper.utility.pool_error import PoolError, PoolErrorCodes
 
 try:
     from epplib.commands import Hello
@@ -33,10 +33,13 @@ class EPPConnectionPool(ConnectionPool):
         try:
             connection = socket.connect()
             return connection
-        except LoginError as err:
+        except Exception as err:
-            message = "_new_connection failed to execute due to a registry login error."
+            message = f"Failed to execute due to a registry login error: {err}"
             logger.error(message, exc_info=True)
-            raise RegistryError(message) from err
+            # We want to raise a pool error rather than a LoginError here
+            # because if this occurs internally, we should handle this
+            # differently than we otherwise would for LoginError.
+            raise PoolError(code=PoolErrorCodes.NEW_CONNECTION_FAILED) from err
 
     def _keepalive(self, c):
         """Sends a command to the server to keep the connection alive."""
@@ -44,8 +47,9 @@ class EPPConnectionPool(ConnectionPool):
             # Sends a ping to EPPLib
             c.send(Hello())
         except Exception as err:
-            logger.error("Failed to keep the connection alive.", exc_info=True)
+            message = "Failed to keep the connection alive."
-            raise RegistryError("Failed to keep the connection alive.") from err
+            logger.error(message, exc_info=True)
+            raise PoolError(code=PoolErrorCodes.KEEP_ALIVE_FAILED) from err
 
     def _create_socket(self, client, login) -> Socket:
         """Creates and returns a socket instance"""
@@ -64,7 +68,6 @@ class EPPConnectionPool(ConnectionPool):
             # Clear the semaphore
             for i in range(self.lock.counter):
                 self.lock.release()
-        # TODO - connection pool err
         except Exception as err:
             logger.error("Could not kill all connections.")
             raise err

src/epplibwrapper/utility/pool_error.py

@@ -0,0 +1,44 @@
+from enum import IntEnum
+
+
+class PoolErrorCodes(IntEnum):
+    """Used in the PoolError class for
+    error mapping.
+
+    Overview of contact error codes:
+        - 2000 KILL_ALL_FAILED
+        - 2001 NEW_CONNECTION_FAILED
+        - 2002 KEEP_ALIVE_FAILED
+    """
+
+    KILL_ALL_FAILED = 2000
+    NEW_CONNECTION_FAILED = 2001
+    KEEP_ALIVE_FAILED = 2002
+
+
+class PoolError(Exception):
+    """
+    Overview of contact error codes:
+        - 2000 KILL_ALL_FAILED
+        - 2001 NEW_CONNECTION_FAILED
+        - 2002 KEEP_ALIVE_FAILED
+    """
+
+    # For linter
+    kill_failed = "Could not kill all connections."
+    conn_failed = "Failed to execute due to a registry login error."
+    alive_failed = "Failed to keep the connection alive."
+    _error_mapping = {
+        PoolErrorCodes.KILL_ALL_FAILED: kill_failed,
+        PoolErrorCodes.NEW_CONNECTION_FAILED: conn_failed,
+        PoolErrorCodes.KEEP_ALIVE_FAILED: alive_failed,
+    }
+
+    def __init__(self, *args, code=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.code = code
+        if self.code in self._error_mapping:
+            self.message = self._error_mapping.get(self.code)
+
+    def __str__(self):
+        return f"{self.message}"

src/registrar/models/domain.py

@@ -16,32 +16,25 @@ from registrar.utility.errors import (
     NameserverErrorCodes as nsErrorCodes,
 )
 
-from registrar.models.utility.contact_error import ContactError, ContactErrorCodes
+from epplibwrapper import (
+    CLIENT as registry,
+    commands,
+    common as epp,
+    extensions,
+    info as eppInfo,
+    RegistryError,
+    ErrorCode,
+)
 
+from registrar.models.utility.contact_error import ContactError, ContactErrorCodes
 
 from .utility.domain_field import DomainField
 from .utility.domain_helper import DomainHelper
 from .utility.time_stamped_model import TimeStampedModel
 
 from .public_contact import PublicContact
+
 logger = logging.getLogger(__name__)
-try:
-    from epplibwrapper import (
-        CLIENT as registry,
-        commands,
-        common as epp,
-        extensions,
-        info as eppInfo,
-        RegistryError,
-        ErrorCode,
-    )
-except ImportError as err:
-    logger.error("An import error occured....")
-    logger.error(err)
-    raise err
-
-
-
 
 
 class Domain(TimeStampedModel, DomainHelper):