From 057e8217fffb95c6d9b2fc1994bf55a148ce0f68 Mon Sep 17 00:00:00 2001
From: Robert Tuck <robert.tuck@diamond.ac.uk>
Date: Fri, 13 Sep 2024 15:07:01 +0100
Subject: [PATCH] (#180) Remove dependency on pkg_resources in favour of
 importlib

---
 src/workflows/services/__init__.py  |  7 ++-----
 src/workflows/transport/__init__.py |  8 ++------
 tests/services/test.py              | 13 +++++++++++++
 tests/transport/test.py             | 14 +++++++++++++-
 4 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/workflows/services/__init__.py b/src/workflows/services/__init__.py
index 9b008ab7..458ad7ea 100644
--- a/src/workflows/services/__init__.py
+++ b/src/workflows/services/__init__.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import pkg_resources
+from importlib.metadata import entry_points
 
 
 def lookup(service: str):
@@ -25,10 +25,7 @@ def get_known_services():
         setattr(
             get_known_services,
             "cache",
-            {
-                e.name: e.load
-                for e in pkg_resources.iter_entry_points("workflows.services")
-            },
+            {e.name: e.load for e in entry_points()["workflows.services"]},
         )
     register = get_known_services.cache.copy()
     return register
diff --git a/src/workflows/transport/__init__.py b/src/workflows/transport/__init__.py
index 0ede6b1b..e62a4e1f 100644
--- a/src/workflows/transport/__init__.py
+++ b/src/workflows/transport/__init__.py
@@ -2,10 +2,9 @@
 
 import argparse
 import optparse
+from importlib.metadata import entry_points
 from typing import TYPE_CHECKING, Type
 
-import pkg_resources
-
 if TYPE_CHECKING:
     from .common_transport import CommonTransport
 
@@ -61,9 +60,6 @@ def get_known_transports() -> dict[str, Type[CommonTransport]]:
         setattr(
             get_known_transports,
             "cache",
-            {
-                e.name: e.load()
-                for e in pkg_resources.iter_entry_points("workflows.transport")
-            },
+            {e.name: e.load() for e in entry_points()["workflows.transport"]},
         )
     return get_known_transports.cache.copy()  # type: ignore
diff --git a/tests/services/test.py b/tests/services/test.py
index 388b4d47..d0277811 100644
--- a/tests/services/test.py
+++ b/tests/services/test.py
@@ -1,8 +1,21 @@
 from __future__ import annotations
 
 import workflows.services
+from workflows.services.common_service import CommonService
 
 
 def test_known_services_is_a_dictionary():
     """Check services register build in CommonService."""
     assert isinstance(workflows.services.get_known_services(), dict)
+
+
+def test_enumerate_services():
+    """Verify we can discover the installed services."""
+    services = workflows.services.get_known_services()
+    assert services.keys() == {
+        "SampleConsumer",
+        "SampleProducer",
+        "SampleTxn",
+        "SampleTxnProducer",
+    }
+    assert all([issubclass(service(), CommonService) for service in services.values()])
diff --git a/tests/transport/test.py b/tests/transport/test.py
index 9c55a1f8..bde7fab2 100644
--- a/tests/transport/test.py
+++ b/tests/transport/test.py
@@ -1,11 +1,23 @@
 from __future__ import annotations
 
 import workflows.transport
+from workflows.transport.common_transport import CommonTransport
 
 
 def test_known_transports_is_a_dictionary():
     """Check transport register build in CommonTransport."""
-    assert isinstance(workflows.transport.get_known_transports(), dict)
+    transports = workflows.transport.get_known_transports()
+    print(transports)
+    assert isinstance(transports, dict)
+
+
+def test_enumerate_transports():
+    """Verify we can discover the installed transports."""
+    transports = workflows.transport.get_known_transports()
+    assert transports.keys() == {"OfflineTransport", "PikaTransport", "StompTransport"}
+    assert all(
+        [issubclass(transport, CommonTransport) for transport in transports.values()]
+    )
 
 
 def test_load_any_transport():