From ecf2fa53e3eeef934fb72ac04ebb428d68fb6640 Mon Sep 17 00:00:00 2001
From: Samuel Cormier-Iijima <samuel@cormier-iijima.com>
Date: Tue, 4 Jun 2024 11:09:36 -0400
Subject: [PATCH] Address review comments, remove redis support

---
 README.md                                   |  1 +
 asyncio_connection_pool/contrib/aioredis.py | 22 --------------
 asyncio_connection_pool/contrib/datadog.py  | 32 ++++++++++++---------
 pyproject.toml                              |  1 -
 test/test_pool.py                           |  2 +-
 5 files changed, 21 insertions(+), 37 deletions(-)
 delete mode 100644 asyncio_connection_pool/contrib/aioredis.py

diff --git a/README.md b/README.md
index 61bc760..3988b41 100644
--- a/README.md
+++ b/README.md
@@ -4,6 +4,7 @@
 [![PyPI](https://img.shields.io/pypi/v/asyncio-connection-pool?style=flat)][package]
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/asyncio-connection-pool?style=flat)][package]
 [![codecov](https://codecov.io/gh/fellowapp/asyncio-connection-pool/graph/badge.svg?token=F3D4D9EG6M)](https://codecov.io/gh/fellowapp/asyncio-connection-pool)
+[![License](https://img.shields.io/pypi/l/prosemirror.svg?style=flat)](https://github.com/fellowapp/asyncio-connection-pool/blob/master/LICENSE.md)
 [![Fellow Careers](https://img.shields.io/badge/fellow.app-hiring-576cf7.svg?style=flat)](https://fellow.app/careers/)
 
 [main CI]: https://github.com/fellowapp/asyncio-connection-pool/actions?query=workflow%3ACI+branch%3Amain
diff --git a/asyncio_connection_pool/contrib/aioredis.py b/asyncio_connection_pool/contrib/aioredis.py
deleted file mode 100644
index 911f597..0000000
--- a/asyncio_connection_pool/contrib/aioredis.py
+++ /dev/null
@@ -1,22 +0,0 @@
-from functools import partial
-
-from redis import asyncio as aioredis
-
-from asyncio_connection_pool import ConnectionStrategy
-
-__all__ = ("RedisConnectionStrategy",)
-
-
-class RedisConnectionStrategy(ConnectionStrategy[aioredis.Redis]):
-    def __init__(self, *args, **kwargs):
-        self._create_redis = partial(aioredis.Redis, *args, **kwargs)
-
-    async def make_connection(self) -> aioredis.Redis:
-        return self._create_redis()
-
-    def connection_is_closed(self, conn):
-        return conn.closed
-
-    async def close_connection(self, conn):
-        conn.close()
-        await conn.wait_closed()
diff --git a/asyncio_connection_pool/contrib/datadog.py b/asyncio_connection_pool/contrib/datadog.py
index a5c96e8..57ac49e 100644
--- a/asyncio_connection_pool/contrib/datadog.py
+++ b/asyncio_connection_pool/contrib/datadog.py
@@ -1,5 +1,5 @@
 from contextlib import AsyncExitStack, asynccontextmanager
-from typing import AsyncIterator, Awaitable, TypeVar
+from typing import Any, AsyncIterator, Awaitable, Coroutine, TypeVar
 
 from datadog.dogstatsd.base import statsd
 from ddtrace import tracer
@@ -85,29 +85,35 @@ def _record_connection_acquiring(self, value: int = 0) -> None:
             tags=self._extra_tags,
         )
 
-    async def _connection_maker(self) -> Conn:
+    def _connection_maker(self) -> Coroutine[Any, Any, Conn]:
         statsd.increment(
             f"{self._service_name}.pool.getting_connection",
             tags=[*self._extra_tags, "method:new"],
         )
 
-        with tracer.trace(
-            f"{self._service_name}.pool._create_new_connection",
-            service=self._service_name,
-        ):
-            return await super()._connection_maker()
+        async def connection_maker(self) -> Conn:
+            with tracer.trace(
+                f"{self._service_name}.pool._create_new_connection",
+                service=self._service_name,
+            ):
+                return await super()._connection_maker()
+
+        return connection_maker(self)
 
-    async def _connection_waiter(self) -> Conn:
+    def _connection_waiter(self) -> Coroutine[Any, Any, Conn]:
         statsd.increment(
             f"{self._service_name}.pool.getting_connection",
             tags=[*self._extra_tags, "method:wait"],
         )
 
-        with tracer.trace(
-            f"{self._service_name}.pool._wait_for_connection",
-            service=self._service_name,
-        ):
-            return await super()._connection_waiter()
+        async def connection_waiter(self) -> Conn:
+            with tracer.trace(
+                f"{self._service_name}.pool._wait_for_connection",
+                service=self._service_name,
+            ):
+                return await super()._connection_waiter()
+
+        return connection_waiter(self)
 
     def _get_conn(self) -> Awaitable[Conn]:
         if not self.available.empty():
diff --git a/pyproject.toml b/pyproject.toml
index 1cd7119..a383b99 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,6 @@ dependencies = []
 
 [project.optional-dependencies]
 datadog = ["ddtrace", "datadog"]
-aioredis = ["redis>=4.2"]
 dev = [
     "codecov~=2.1",
     "mypy~=1.9",
diff --git a/test/test_pool.py b/test/test_pool.py
index d6f468b..915b0bd 100644
--- a/test/test_pool.py
+++ b/test/test_pool.py
@@ -131,7 +131,7 @@ async def worker():
     assert pool.in_use == nworkers, "all workers should have their connections now"
     ev2.set()
     await asyncio.gather(*coros)
-    assert pool.currently_allocating == 0
+    assert pool.in_use == 0
     assert (
         pool.available.qsize() == nworkers
     ), "all workers should have returned their connections"