zydronium/manage.get.gov
1
0
Fork 0
mirror of https://github.com/cisagov/manage.get.gov.git synced 2025-05-18 18:39:21 +02:00
Code Issues Projects Releases Packages Wiki Activity

Tests / code cleanup

Browse source
This commit is contained in:
zandercymatics 2023-10-17 11:21:18 -06:00
parent 3aac1e38df
commit d7c7fb9d23
No known key found for this signature in database
GPG key ID: FF4636ABEC9682B7
6 changed files with 188 additions and 130 deletions
Download patch file Download diff file Expand all files Collapse all files

42
src/epplibwrapper/client.py
View file

@ -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,6 +80,7 @@ class EPPLibWrapper:
# Tracks the status of the pool
self.pool_status = PoolStatus()
if start_connection_pool:
self.start_connection_pool()
def _send(self, command):
@ -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(
"Can't contact the Registry. Please try again later"
)
logger.error("Can't contact the Registry. Pool was not running.")
raise RegistryError("Can't contact the Registry. Pool was not running.")
except RegistryError as err:
raise err
finally:
@ -159,23 +160,30 @@ 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
@ -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(

4
src/epplibwrapper/socket.py
View file

@ -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)

158
src/epplibwrapper/tests/test_pool.py
View file

@ -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(MockEppLib):
class TestConnectionPool(TestCase):
"""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):
super().tearDown()
# Mock a successful connection
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):
# setup
session = self.client.session
session["state"] = "TEST" # nosec B105
session.save()
# mock
@patch("djangooidc.views.CLIENT", autospec=True)
def fake_pool(self, mock_client):
# mock client
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 = client._pool
# 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=[],
pool = EPPConnectionPool(
client=mock_client, login=mock_login, options=self.pool_options
)
# 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)
return pool
@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"""
raise
"""A .send is invoked on the pool successfully"""
# 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"])

17
src/epplibwrapper/utility/pool.py
View file

@ -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:
message = "_new_connection failed to execute due to a registry login error."
except Exception as err:
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)
raise RegistryError("Failed to keep the connection alive.") from err
message = "Failed to keep the connection alive."
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

44
src/epplibwrapper/utility/pool_error.py Normal file
View file

@ -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}"

21
src/registrar/models/domain.py
View file

@ -16,16 +16,6 @@ from registrar.utility.errors import (
NameserverErrorCodes as nsErrorCodes,
)
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,
@ -35,13 +25,16 @@ try:
RegistryError,
ErrorCode,
)
except ImportError as err:
logger.error("An import error occured....")
logger.error(err)
raise err
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__)
class Domain(TimeStampedModel, DomainHelper):