refactor: parametrize unit tests

Author: sophia-bq

aws_advanced_python_wrapper/read_write_splitting_plugin.py

@@ -50,6 +50,7 @@ class ReadWriteSplittingConnectionManager(Plugin):
         "Connection.set_read_only",
     }
     _POOL_PROVIDER_CLASS_NAME = "aws_advanced_python_wrapper.sql_alchemy_connection_provider.SqlAlchemyPooledConnectionProvider"
+    _CLOSE_METHOD = "Connection.close"
 
     def __init__(
         self,
@@ -94,8 +95,7 @@ def connect(
         connect_func: Callable,
     ) -> Connection:
         return self._connection_handler.get_verified_initial_connection(
-            host_info, props, is_initial_connection, connect_func
-        )
+            host_info, is_initial_connection, lambda x: self._plugin_service.connect(x, props, self), connect_func)
 
     def notify_connection_changed(
         self, changes: Set[ConnectionEvent]
@@ -140,13 +140,13 @@ def execute(
             if isinstance(ex, FailoverError):
                 logger.debug(
                     "ReadWriteSplittingPlugin.FailoverExceptionWhileExecutingCommand",
-                    method_name,
+                    method_name
                 )
                 self._close_idle_connections()
             else:
                 logger.debug(
                     "ReadWriteSplittingPlugin.ExceptionWhileExecutingCommand",
-                    method_name,
+                    method_name
                 )
             raise ex
 
@@ -184,13 +184,13 @@ def _set_reader_connection(
         )
 
     def _initialize_writer_connection(self):
-        conn, writer_host = self._connection_handler.open_new_writer_connection()
+        conn, writer_host = self._connection_handler.open_new_writer_connection(lambda x: self._plugin_service.connect(x, self._properties, self))
 
         if conn is None:
             self.log_and_raise_exception(
                 "ReadWriteSplittingPlugin.FailedToConnectToWriter"
             )
-            return
+            return None
 
         provider = self._conn_provider_manager.get_connection_provider(
             writer_host, self._properties
@@ -335,9 +335,7 @@ def _switch_to_reader_connection(self):
                     self._reader_host_info.url,
                 )
 
-                ReadWriteSplittingConnectionManager.close_connection(
-                    self._reader_connection
-                )
+                ReadWriteSplittingConnectionManager.close_connection(self._reader_connection, driver_dialect)
                 self._reader_connection = None
                 self._reader_host_info = None
                 self._initialize_reader_connection()
@@ -356,7 +354,7 @@ def _initialize_reader_connection(self):
             )
             return
 
-        conn, reader_host = self._connection_handler.open_new_reader_connection()
+        conn, reader_host = self._connection_handler.open_new_reader_connection(lambda x: self._plugin_service.connect(x, self._properties, self))
 
         if conn is None or reader_host is None:
             self.log_and_raise_exception("ReadWriteSplittingPlugin.NoReadersAvailable")
@@ -392,7 +390,7 @@ def _close_connection_if_idle(self, internal_conn: Optional[Connection]):
             if internal_conn != current_conn and self._is_connection_usable(
                 internal_conn, driver_dialect
             ):
-                internal_conn.close()
+                driver_dialect.execute(ReadWriteSplittingConnectionManager._CLOSE_METHOD, lambda: internal_conn.close())
                 if internal_conn == self._writer_connection:
                     self._writer_connection = None
                     self._writer_host_info = None
@@ -420,10 +418,8 @@ def log_and_raise_exception(log_msg: str):
         raise ReadWriteSplittingError(Messages.get(log_msg))
 
     @staticmethod
-    def _is_connection_usable(
-        conn: Optional[Connection], driver_dialect: Optional[DriverDialect]
-    ):
-        if conn is None or driver_dialect is None:
+    def _is_connection_usable(conn: Optional[Connection], driver_dialect: DriverDialect):
+        if conn is None:
             return False
         try:
             return not driver_dialect.is_closed(conn)
@@ -432,10 +428,10 @@ def _is_connection_usable(
             return False
 
     @staticmethod
-    def close_connection(connection: Optional[Connection]):
-        if connection is not None:
+    def close_connection(conn: Optional[Connection], driver_dialect: DriverDialect):
+        if conn is not None:
             try:
-                connection.close()
+                driver_dialect.execute(ReadWriteSplittingConnectionManager._CLOSE_METHOD, lambda: conn.close())
             except Exception:
                 # Swallow exception
                 return
@@ -456,21 +452,23 @@ def host_list_provider_service(self, new_value: int) -> None:
 
     def open_new_writer_connection(
         self,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
     ) -> tuple[Optional[Connection], Optional[HostInfo]]:
         """Open a writer connection."""
         ...
 
     def open_new_reader_connection(
         self,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
     ) -> tuple[Optional[Connection], Optional[HostInfo]]:
         """Open a reader connection."""
         ...
 
     def get_verified_initial_connection(
         self,
         host_info: HostInfo,
-        props: Properties,
         is_initial_connection: bool,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
         connect_func: Callable,
     ) -> Connection:
         """Verify initial connection or return normal workflow."""
@@ -516,9 +514,8 @@ class TopologyBasedConnectionHandler(ConnectionHandler):
 
     def __init__(self, plugin_service: PluginService, props: Properties):
         self._plugin_service: PluginService = plugin_service
-        self._properties: Properties = props
         self._host_list_provider_service: Optional[HostListProviderService] = None
-        strategy = WrapperProperties.READER_HOST_SELECTOR_STRATEGY.get(self._properties)
+        strategy = WrapperProperties.READER_HOST_SELECTOR_STRATEGY.get(props)
         if strategy is not None:
             self._reader_selector_strategy = strategy
         else:
@@ -539,17 +536,19 @@ def host_list_provider_service(self, new_value: HostListProviderService) -> None
 
     def open_new_writer_connection(
         self,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
     ) -> tuple[Optional[Connection], Optional[HostInfo]]:
         writer_host = self._get_writer()
         if writer_host is None:
             return None, None
 
-        conn = self._plugin_service.connect(writer_host, self._properties, None)
+        conn = plugin_service_connect_func(writer_host)
 
         return conn, writer_host
 
     def open_new_reader_connection(
         self,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
     ) -> tuple[Optional[Connection], Optional[HostInfo]]:
         conn: Optional[Connection] = None
         reader_host: Optional[HostInfo] = None
@@ -561,7 +560,7 @@ def open_new_reader_connection(
             )
             if host is not None:
                 try:
-                    conn = self._plugin_service.connect(host, self._properties, None)
+                    conn = plugin_service_connect_func(host)
                     reader_host = host
                     break
                 except Exception:
@@ -574,8 +573,8 @@ def open_new_reader_connection(
     def get_verified_initial_connection(
         self,
         host_info: HostInfo,
-        props: Properties,
         is_initial_connection: bool,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
         connect_func: Callable,
     ) -> Connection:
         if not self._plugin_service.accepts_strategy(
@@ -670,12 +669,9 @@ def _get_writer(self) -> Optional[HostInfo]:
 
 
 class ReadWriteSplittingPlugin(ReadWriteSplittingConnectionManager):
-    def __init__(self, plugin_service, props: Properties):
+    def __init__(self, plugin_service: PluginService, props: Properties):
         # The read/write splitting plugin handles connections based on topology.
-        connection_handler = TopologyBasedConnectionHandler(
-            plugin_service,
-            props,
-        )
+        connection_handler = TopologyBasedConnectionHandler(plugin_service, props)
 
         super().__init__(plugin_service, props, connection_handler)
 

aws_advanced_python_wrapper/simple_read_write_splitting_plugin.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from time import perf_counter_ns, sleep
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Type
+from typing import TYPE_CHECKING, Callable, Optional, Type, TypeVar
 
 from aws_advanced_python_wrapper.host_availability import HostAvailability
 from aws_advanced_python_wrapper.read_write_splitting_plugin import (
@@ -67,15 +67,10 @@ def __init__(self, plugin_service: PluginService, props: Properties):
             )
 
         self._plugin_service: PluginService = plugin_service
-        self._properties: Properties = props
         self._rds_utils: RdsUtils = RdsUtils()
         self._host_list_provider_service: Optional[HostListProviderService] = None
-        self._write_endpoint_host_info: HostInfo = self._create_host_info(
-            self._write_endpoint, HostRole.WRITER
-        )
-        self._read_endpoint_host_info: HostInfo = self._create_host_info(
-            self._read_endpoint, HostRole.READER
-        )
+        self._write_endpoint_host_info: HostInfo = self._create_host_info(self._write_endpoint, HostRole.WRITER)
+        self._read_endpoint_host_info: HostInfo = self._create_host_info(self._read_endpoint, HostRole.READER)
 
     @property
     def host_list_provider_service(self) -> Optional[HostListProviderService]:
@@ -87,39 +82,29 @@ def host_list_provider_service(self, new_value: HostListProviderService) -> None
 
     def open_new_writer_connection(
         self,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
     ) -> tuple[Optional[Connection], Optional[HostInfo]]:
-        conn: Optional[Connection] = None
         if self._verify_new_connections:
-            conn = self._get_verified_connection(
-                self._properties, self._write_endpoint_host_info, HostRole.WRITER
-            )
-        else:
-            conn = self._plugin_service.connect(
-                self._write_endpoint_host_info, self._properties, None
-            )
+            return self._get_verified_connection(self._write_endpoint_host_info, HostRole.WRITER, plugin_service_connect_func), \
+                self._write_endpoint_host_info
 
-        return conn, self._write_endpoint_host_info
+        return plugin_service_connect_func(self._write_endpoint_host_info), self._write_endpoint_host_info
 
     def open_new_reader_connection(
         self,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
     ) -> tuple[Optional[Connection], Optional[HostInfo]]:
-        conn: Optional[Connection] = None
         if self._verify_new_connections:
-            conn = self._get_verified_connection(
-                self._properties, self._read_endpoint_host_info, HostRole.READER
-            )
-        else:
-            conn = self._plugin_service.connect(
-                self._read_endpoint_host_info, self._properties, None
-            )
+            return self._get_verified_connection(self._read_endpoint_host_info, HostRole.READER, plugin_service_connect_func), \
+                self._read_endpoint_host_info
 
-        return conn, self._read_endpoint_host_info
+        return plugin_service_connect_func(self._read_endpoint_host_info), self._read_endpoint_host_info
 
     def get_verified_initial_connection(
         self,
         host_info: HostInfo,
-        props: Properties,
         is_initial_connection: bool,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
         connect_func: Callable,
     ) -> Connection:
         if not is_initial_connection or not self._verify_new_connections:
@@ -133,34 +118,30 @@ def get_verified_initial_connection(
             url_type == RdsUrlType.RDS_WRITER_CLUSTER
             or self._verify_opened_connection_type == HostRole.WRITER
         ):
-            conn = self._get_verified_connection(
-                props, host_info, HostRole.WRITER, connect_func
-            )
+            conn = self._get_verified_connection(host_info, HostRole.WRITER, plugin_service_connect_func, connect_func)
         elif (
             url_type == RdsUrlType.RDS_READER_CLUSTER
             or self._verify_opened_connection_type == HostRole.READER
         ):
-            conn = self._get_verified_connection(
-                props, host_info, HostRole.READER, connect_func
-            )
+            conn = self._get_verified_connection(host_info, HostRole.READER, plugin_service_connect_func, connect_func)
 
         if conn is None:
             conn = connect_func()
 
-        self._set_initial_connection_host_info(conn, host_info)
+        self._set_initial_connection_host_info(host_info)
         return conn
 
-    def _set_initial_connection_host_info(self, conn: Connection, host_info: HostInfo):
+    def _set_initial_connection_host_info(self, host_info: HostInfo):
         if self._host_list_provider_service is None:
             return
 
         self._host_list_provider_service.initial_connection_host_info = host_info
 
     def _get_verified_connection(
         self,
-        props: Properties,
         host_info: HostInfo,
         role: HostRole,
+        plugin_service_connect_func: Callable[[HostInfo], Connection],
         connect_func: Optional[Callable] = None,
     ) -> Optional[Connection]:
         end_time_nano = perf_counter_ns() + (self._connect_retry_timeout_ms * 1000000)
@@ -174,9 +155,7 @@ def _get_verified_connection(
                 if connect_func is not None:
                     candidate_conn = connect_func()
                 elif host_info is not None:
-                    candidate_conn = self._plugin_service.connect(
-                        host_info, props, None
-                    )
+                    candidate_conn = plugin_service_connect_func(host_info)
                 else:
                     return None
 
@@ -187,14 +166,14 @@ def _get_verified_connection(
                 actual_role = self._plugin_service.get_host_role(candidate_conn)
 
                 if actual_role != role:
-                    ReadWriteSplittingConnectionManager.close_connection(candidate_conn)
+                    ReadWriteSplittingConnectionManager.close_connection(candidate_conn, self._plugin_service.driver_dialect)
                     self._delay()
                     continue
 
                 return candidate_conn
 
             except Exception:
-                ReadWriteSplittingConnectionManager.close_connection(candidate_conn)
+                ReadWriteSplittingConnectionManager.close_connection(candidate_conn, self._plugin_service.driver_dialect)
                 self._delay()
 
         return None
@@ -249,28 +228,18 @@ def is_reader_host(self, current_host: HostInfo) -> bool:
             or current_host.url.casefold() == self._read_endpoint.casefold()
         )
 
-    def _create_host_info(self, endpoint, role: HostRole) -> HostInfo:
+    def _create_host_info(self, endpoint: str, role: HostRole) -> HostInfo:
         endpoint = endpoint.strip()
         host = endpoint
-        port = self._plugin_service.database_dialect.default_port
+        port = self._plugin_service.database_dialect.default_port if not self._plugin_service.current_host_info.is_port_specified() \
+            else self._plugin_service.current_host_info.port
         colon_index = endpoint.rfind(":")
 
         if colon_index != -1:
+            host = endpoint[:colon_index]
             port_str = endpoint[colon_index + 1:]
             if port_str.isdigit():
-                host = endpoint[:colon_index]
                 port = int(port_str)
-            else:
-                if (
-                    self._host_list_provider_service is not None
-                    and self._host_list_provider_service.initial_connection_host_info
-                    is not None
-                    and self._host_list_provider_service.initial_connection_host_info.port
-                    != HostInfo.NO_PORT
-                ):
-                    port = (
-                        self._host_list_provider_service.initial_connection_host_info.port
-                    )
 
         return HostInfo(
             host=host, port=port, role=role, availability=HostAvailability.AVAILABLE
@@ -322,12 +291,9 @@ def _parse_connection_type(phase_str: Optional[str]) -> HostRole:
 
 
 class SimpleReadWriteSplittingPlugin(ReadWriteSplittingConnectionManager):
-    def __init__(self, plugin_service, props: Properties):
+    def __init__(self, plugin_service: PluginService, props: Properties):
         # The simple read/write splitting plugin handles connections based on configuration parameter endpoints.
-        connection_handler = EndpointBasedConnectionHandler(
-            plugin_service,
-            props,
-        )
+        connection_handler = EndpointBasedConnectionHandler(plugin_service, props)
 
         super().__init__(plugin_service, props, connection_handler)