From 0dde277c864d308075042dd74a740b86884c191c Mon Sep 17 00:00:00 2001
From: zandercymatics <141044360+zandercymatics@users.noreply.github.com>
Date: Tue, 17 Oct 2023 12:40:15 -0600
Subject: [PATCH] Test changes

---
 src/epplibwrapper/client.py          | 18 +++++++++++++++---
 src/epplibwrapper/tests/test_pool.py | 17 +++++++++++++----
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/epplibwrapper/client.py b/src/epplibwrapper/client.py
index 920dc284d..680a6661c 100644
--- a/src/epplibwrapper/client.py
+++ b/src/epplibwrapper/client.py
@@ -121,6 +121,7 @@ class EPPLibWrapper:
             logger.error(message, exc_info=True)
             raise RegistryError(message) from err
         else:
+            print(f"test thing {response}")
             if response.code >= 2000:
                 raise RegistryError(response.msg, code=response.code)
             else:
@@ -160,6 +161,16 @@ class EPPLibWrapper:
                 else:  # don't try again
                     raise err
 
+    def get_pool(self):
+        """Get the current pool instance"""
+        return self._pool
+    
+    def _create_pool(self, client, login, options):
+        """Creates and returns new pool instance"""
+        return EPPConnectionPool(
+                client, login, options
+        )
+
     def start_connection_pool(
         self, restart_pool_if_exists=True, try_start_if_invalid=False
     ):
@@ -193,9 +204,10 @@ class EPPLibWrapper:
                 logger.info("Connection pool restarting...")
                 self.kill_pool()
 
-            self._pool = EPPConnectionPool(
-                client=self._client, login=self._login, options=self.pool_options
-            )
+            self._pool = self._create_pool(
+                self._client, self._login, self.pool_options
+            ) 
+            
             self.pool_status.pool_running = True
             self.pool_status.pool_hanging = False
 
diff --git a/src/epplibwrapper/tests/test_pool.py b/src/epplibwrapper/tests/test_pool.py
index 6573a64f5..e62135221 100644
--- a/src/epplibwrapper/tests/test_pool.py
+++ b/src/epplibwrapper/tests/test_pool.py
@@ -3,9 +3,11 @@ from unittest.mock import MagicMock, patch
 from django.conf import settings
 
 from django.test import TestCase
+from epplibwrapper.client import EPPLibWrapper
 from epplibwrapper.utility.pool import EPPConnectionPool
 from registrar.models.domain import Domain
 from registrar.models.domain import registry
+from contextlib import ExitStack
 
 import logging
 
@@ -120,16 +122,23 @@ class TestConnectionPool(TestCase):
         # Fake data for the _pool object
         domain, _ = Domain.objects.get_or_create(name="freeman.gov")
 
+        def start_fake_connection(self):
+            registry.pool_status.pool_running = True
+            registry.pool_status.connection_success = True
+            registry._pool = registry.get_pool()
+        
         # 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)
-
+        with ExitStack() as stack:
+            stack.enter_context(patch.object(EPPLibWrapper, "get_pool", self.fake_pool))
+            stack.enter_context(patch.object(EPPLibWrapper, "start_connection_pool", start_fake_connection))
+            expected_contact = domain.security_contact
+        
         # 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"])
+        self.assertEqual(len(registry.get_pool().conn), self.pool_options["size"])