Close all and counts

adafruit_connection_manager.py

@@ -99,7 +99,8 @@ def create_fake_ssl_context(
  return _FakeSSLContext(iface)
 
 
-_global_socketpool = {}
+_global_connection_managers = {}
+_global_socketpools = {}
 _global_ssl_contexts = {}
 
 
@@ -113,7 +114,7 @@ def get_radio_socketpool(radio):
  * Using a WIZ5500 (Like the Adafruit Ethernet FeatherWing)
  """
  class_name = radio.__class__.__name__
- if class_name not in _global_socketpool:
+ if class_name not in _global_socketpools:
  if class_name == "Radio":
  import ssl # pylint: disable=import-outside-toplevel
 
@@ -151,10 +152,10 @@ def get_radio_socketpool(radio):
  else:
  raise AttributeError(f"Unsupported radio class: {class_name}")
 
- _global_socketpool[class_name] = pool
+ _global_socketpools[class_name] = pool
  _global_ssl_contexts[class_name] = ssl_context
 
- return _global_socketpool[class_name]
+ return _global_socketpools[class_name]
 
 
 def get_radio_ssl_context(radio):
@@ -186,10 +187,10 @@ def __init__(
  self._available_socket = {}
  self._open_sockets = {}
 
- def _free_sockets(self) -> None:
+ def _free_sockets(self, force: bool = False) -> None:
  available_sockets = []
  for socket, free in self._available_socket.items():
- if free:
+ if free or force:
  available_sockets.append(socket)
 
  for socket in available_sockets:
@@ -203,6 +204,18 @@ def _get_key_for_socket(self, socket):
  except StopIteration:
  return None
 
+ @property
+ def open_sockets(self) -> int:
+ """Get the count of open sockets"""
+ return len(self._open_sockets)
+
+ @property
+ def freeable_open_sockets(self) -> int:
+ """Get the count of freeable open sockets"""
+ return len(
+ [socket for socket, free in self._available_socket.items() if free is True]
+ )
+
  def close_socket(self, socket: SocketType) -> None:
  """Close a previously opened socket."""
  if socket not in self._open_sockets.values():
@@ -306,11 +319,42 @@ def get_socket(
 # global helpers
 
 
-_global_connection_manager = {}
+def connection_manager_close_all(
+ socket_pool: Optional[SocketpoolModuleType] = None, release_references: bool = False
+) -> None:
+ """Close all open sockets for pool"""
+ if socket_pool:
+ socket_pools = [socket_pool]
+ else:
+ socket_pools = _global_connection_managers.keys()
+
+ for pool in socket_pools:
+ connection_manager = _global_connection_managers.get(pool, None)
+ if connection_manager is None:
+ raise RuntimeError("SocketPool not managed")
+
+ connection_manager._free_sockets(force=True) # pylint: disable=protected-access
+
+ if release_references:
+ radio_key = None
+ for radio_check, pool_check in _global_socketpools.items():
+ if pool == pool_check:
+ radio_key = radio_check
+ break
+
+ if radio_key:
+ if radio_key in _global_socketpools:
+ del _global_socketpools[radio_key]
+
+ if radio_key in _global_ssl_contexts:
+ del _global_ssl_contexts[radio_key]
+
+ if pool in _global_connection_managers:
+ del _global_connection_managers[pool]
 
 
 def get_connection_manager(socket_pool: SocketpoolModuleType) -> ConnectionManager:
  """Get the ConnectionManager singleton for the given pool"""
- if socket_pool not in _global_connection_manager:
- _global_connection_manager[socket_pool] = ConnectionManager(socket_pool)
- return _global_connection_manager[socket_pool]
+ if socket_pool not in _global_connection_managers:
+ _global_connection_managers[socket_pool] = ConnectionManager(socket_pool)
+ return _global_connection_managers[socket_pool]

examples/connectionmanager_helpers.py

@@ -24,14 +24,38 @@
 
 # get request session
 requests = adafruit_requests.Session(pool, ssl_context)
+connection_manager = adafruit_connection_manager.get_connection_manager(pool)
+print("-" * 40)
+print("Nothing yet opened")
+print(f"Open Sockets: {connection_manager.open_sockets}")
+print(f"Freeable Open Sockets: {connection_manager.freeable_open_sockets}")
 
 # make request
 print("-" * 40)
-print(f"Fetching from {TEXT_URL}")
+print(f"Fetching from {TEXT_URL} in a context handler")
+with requests.get(TEXT_URL) as response:
+ response_text = response.text
+ print(f"Text Response {response_text}")
+
+print("-" * 40)
+print("1 request, opened and freed")
+print(f"Open Sockets: {connection_manager.open_sockets}")
+print(f"Freeable Open Sockets: {connection_manager.freeable_open_sockets}")
 
+print("-" * 40)
+print(f"Fetching from {TEXT_URL} not in a context handler")
 response = requests.get(TEXT_URL)
-response_text = response.text
-response.close()
 
-print(f"Text Response {response_text}")
 print("-" * 40)
+print("1 request, opened but not freed")
+print(f"Open Sockets: {connection_manager.open_sockets}")
+print(f"Freeable Open Sockets: {connection_manager.freeable_open_sockets}")
+
+print("-" * 40)
+print("Closing everything in the pool")
+adafruit_connection_manager.connection_manager_close_all(pool)
+
+print("-" * 40)
+print("Everything closed")
+print(f"Open Sockets: {connection_manager.open_sockets}")
+print(f"Freeable Open Sockets: {connection_manager.freeable_open_sockets}")

tests/conftest.py

@@ -65,7 +65,11 @@ def adafruit_wiznet5k_with_ssl_socket_module():
 @pytest.fixture(autouse=True)
 def reset_connection_manager(monkeypatch):
  monkeypatch.setattr(
- "adafruit_connection_manager._global_socketpool",
+ "adafruit_connection_manager._global_connection_managers",
+ {},
+ )
+ monkeypatch.setattr(
+ "adafruit_connection_manager._global_socketpools",
  {},
  )
  monkeypatch.setattr(

tests/connection_manager_close_all_test.py

@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: 2024 Justin Myers for Adafruit Industries
+#
+# SPDX-License-Identifier: Unlicense
+
+""" Get Connection Manager Tests """
+
+import mocket
+import pytest
+
+import adafruit_connection_manager
+
+
+def test_connection_manager_close_all_all():
+ mock_pool_1 = mocket.MocketPool()
+ mock_pool_2 = mocket.MocketPool()
+ assert mock_pool_1 != mock_pool_2
+
+ connection_manager_1 = adafruit_connection_manager.get_connection_manager(
+ mock_pool_1
+ )
+ assert connection_manager_1.open_sockets == 0
+ assert connection_manager_1.freeable_open_sockets == 0
+ connection_manager_2 = adafruit_connection_manager.get_connection_manager(
+ mock_pool_2
+ )
+ assert connection_manager_2.open_sockets == 0
+ assert connection_manager_2.freeable_open_sockets == 0
+ assert len(adafruit_connection_manager._global_connection_managers) == 2
+
+ socket_1 = connection_manager_1.get_socket(mocket.MOCK_HOST_1, 80, "http:")
+ assert connection_manager_1.open_sockets == 1
+ assert connection_manager_1.freeable_open_sockets == 0
+ assert connection_manager_2.open_sockets == 0
+ assert connection_manager_2.freeable_open_sockets == 0
+ socket_2 = connection_manager_2.get_socket(mocket.MOCK_HOST_1, 80, "http:")
+ assert connection_manager_2.open_sockets == 1
+ assert connection_manager_2.freeable_open_sockets == 0
+
+ adafruit_connection_manager.connection_manager_close_all()
+ assert connection_manager_1.open_sockets == 0
+ assert connection_manager_1.freeable_open_sockets == 0
+ assert connection_manager_2.open_sockets == 0
+ assert connection_manager_2.freeable_open_sockets == 0
+ socket_1.close.assert_called_once()
+ socket_2.close.assert_called_once()
+
+
+def test_connection_manager_close_all_single():
+ mock_pool_1 = mocket.MocketPool()
+ mock_pool_2 = mocket.MocketPool()
+ assert mock_pool_1 != mock_pool_2
+
+ connection_manager_1 = adafruit_connection_manager.get_connection_manager(
+ mock_pool_1
+ )
+ assert connection_manager_1.open_sockets == 0
+ assert connection_manager_1.freeable_open_sockets == 0
+ connection_manager_2 = adafruit_connection_manager.get_connection_manager(
+ mock_pool_2
+ )
+ assert connection_manager_2.open_sockets == 0
+ assert connection_manager_2.freeable_open_sockets == 0
+ assert len(adafruit_connection_manager._global_connection_managers) == 2
+
+ socket_1 = connection_manager_1.get_socket(mocket.MOCK_HOST_1, 80, "http:")
+ assert connection_manager_1.open_sockets == 1
+ assert connection_manager_1.freeable_open_sockets == 0
+ assert connection_manager_2.open_sockets == 0
+ assert connection_manager_2.freeable_open_sockets == 0
+ socket_2 = connection_manager_2.get_socket(mocket.MOCK_HOST_1, 80, "http:")
+ assert connection_manager_2.open_sockets == 1
+ assert connection_manager_2.freeable_open_sockets == 0
+
+ adafruit_connection_manager.connection_manager_close_all(mock_pool_1)
+ assert connection_manager_1.open_sockets == 0
+ assert connection_manager_1.freeable_open_sockets == 0
+ assert connection_manager_2.open_sockets == 1
+ assert connection_manager_2.freeable_open_sockets == 0
+ socket_1.close.assert_called_once()
+ socket_2.close.assert_not_called()
+
+
+def test_connection_manager_close_all_untracked():
+ mock_pool_1 = mocket.MocketPool()
+ with pytest.raises(RuntimeError) as context:
+ adafruit_connection_manager.connection_manager_close_all(mock_pool_1)
+ assert "SocketPool not managed" in str(context)
+
+
+def test_connection_manager_close_all_single_release_references_false( # pylint: disable=unused-argument
+ circuitpython_socketpool_module, adafruit_esp32spi_socket_module
+):
+ radio_wifi = mocket.MockRadio.Radio()
+ radio_esp = mocket.MockRadio.ESP_SPIcontrol()
+
+ socket_pool_wifi = adafruit_connection_manager.get_radio_socketpool(radio_wifi)
+ socket_pool_esp = adafruit_connection_manager.get_radio_socketpool(radio_esp)
+
+ ssl_context_wifi = adafruit_connection_manager.get_radio_ssl_context(radio_wifi)
+ ssl_context_esp = adafruit_connection_manager.get_radio_ssl_context(radio_esp)
+
+ connection_manager_wifi = adafruit_connection_manager.get_connection_manager(
+ socket_pool_wifi
+ )
+ connection_manager_esp = adafruit_connection_manager.get_connection_manager(
+ socket_pool_esp
+ )
+
+ assert socket_pool_wifi != socket_pool_esp
+ assert ssl_context_wifi != ssl_context_esp
+ assert connection_manager_wifi != connection_manager_esp
+
+ adafruit_connection_manager.connection_manager_close_all(
+ socket_pool_wifi, release_references=False
+ )
+
+ assert socket_pool_wifi in adafruit_connection_manager._global_socketpools.values()
+ assert socket_pool_esp in adafruit_connection_manager._global_socketpools.values()
+
+ assert ssl_context_wifi in adafruit_connection_manager._global_ssl_contexts.values()
+ assert ssl_context_esp in adafruit_connection_manager._global_ssl_contexts.values()
+
+ assert (
+ socket_pool_wifi
+ in adafruit_connection_manager._global_connection_managers.keys()
+ )
+ assert (
+ socket_pool_esp
+ in adafruit_connection_manager._global_connection_managers.keys()
+ )
+
+
+def test_connection_manager_close_all_single_release_references_true( # pylint: disable=unused-argument
+ circuitpython_socketpool_module, adafruit_esp32spi_socket_module
+):
+ radio_wifi = mocket.MockRadio.Radio()
+ radio_esp = mocket.MockRadio.ESP_SPIcontrol()
+
+ socket_pool_wifi = adafruit_connection_manager.get_radio_socketpool(radio_wifi)
+ socket_pool_esp = adafruit_connection_manager.get_radio_socketpool(radio_esp)
+
+ ssl_context_wifi = adafruit_connection_manager.get_radio_ssl_context(radio_wifi)
+ ssl_context_esp = adafruit_connection_manager.get_radio_ssl_context(radio_esp)
+
+ connection_manager_wifi = adafruit_connection_manager.get_connection_manager(
+ socket_pool_wifi
+ )
+ connection_manager_esp = adafruit_connection_manager.get_connection_manager(
+ socket_pool_esp
+ )
+
+ assert socket_pool_wifi != socket_pool_esp
+ assert ssl_context_wifi != ssl_context_esp
+ assert connection_manager_wifi != connection_manager_esp
+
+ adafruit_connection_manager.connection_manager_close_all(
+ socket_pool_wifi, release_references=True
+ )
+
+ assert (
+ socket_pool_wifi not in adafruit_connection_manager._global_socketpools.values()
+ )
+ assert socket_pool_esp in adafruit_connection_manager._global_socketpools.values()
+
+ assert (
+ ssl_context_wifi
+ not in adafruit_connection_manager._global_ssl_contexts.values()
+ )
+ assert ssl_context_esp in adafruit_connection_manager._global_ssl_contexts.values()
+
+ assert (
+ socket_pool_wifi
+ not in adafruit_connection_manager._global_connection_managers.keys()
+ )
+ assert (
+ socket_pool_esp
+ in adafruit_connection_manager._global_connection_managers.keys()
+ )