Removed circular references to LabThing

Author: Joel Collins

labthings/server/default_views/root.py

@@ -0,0 +1,7 @@
+from ..find import current_labthing
+from ..view import View
+
+
+class RootView(View):
+    def get(self):
+        return current_labthing().thing_description.to_dict()

labthings/server/default_views/sockets.py

@@ -0,0 +1,16 @@
+from ..sockets import SocketSubscriber, socket_handler_loop
+from ..find import current_labthing
+
+import logging
+
+
+def socket_handler(ws):
+    # Create a socket subscriber
+    wssub = SocketSubscriber(ws)
+    current_labthing().subscribers.add(wssub)
+    logging.info(f"Added subscriber {wssub}")
+    # Start the socket connection handler loop
+    socket_handler_loop(ws)
+    # Remove the subscriber once the loop returns
+    current_labthing().subscribers.remove(wssub)
+    logging.info(f"Removed subscriber {wssub}")

labthings/server/find.py

@@ -1,5 +1,6 @@
 import logging
 from flask import current_app
+import weakref
 
 from .names import EXTENSION_NAME
 
@@ -13,14 +14,15 @@ def current_labthing(app=None):
     # reach the Flask app object. Just using current_app returns
     # a wrapper, which breaks it's use in Task threads
     if not app:
-        app = current_app._get_current_object()  # skipcq: PYL-W0212
-    if not app:
-        return None
-    logging.debug("Active app extensions:")
-    logging.debug(app.extensions)
-    logging.debug("Active labthing:")
-    logging.debug(app.extensions.get(EXTENSION_NAME))
-    return app.extensions.get(EXTENSION_NAME, None)
+        try:
+            app = current_app._get_current_object()  # skipcq: PYL-W0212
+        except RuntimeError:
+            return None
+    ext = app.extensions.get(EXTENSION_NAME, None)
+    if isinstance(ext, weakref.ref):
+        return ext()
+    else:
+        return ext
 
 
 def registered_extensions(labthing_instance=None):

labthings/server/labthing.py

@@ -17,14 +17,17 @@
 from .spec.utilities import get_spec
 from .spec.td import ThingDescription
 from .decorators import tag
-from .sockets import Sockets, SocketSubscriber, socket_handler_loop
+from .sockets import Sockets
 
 from .default_views.extensions import ExtensionList
 from .default_views.tasks import TaskList, TaskView
 from .default_views.docs import docs_blueprint, SwaggerUIView
+from .default_views.root import RootView
+from .default_views.sockets import socket_handler
 
 from ..core.utilities import get_docstring
 
+import weakref
 import logging
 
 
@@ -72,7 +75,7 @@ def __init__(
 
         # Logging handler
         # TODO: Add cleanup code
-        self.log_handler = LabThingLogger(self)
+        self.log_handler = LabThingLogger()
         logging.getLogger().addHandler(self.log_handler)
 
         self.spec = APISpec(
@@ -87,6 +90,12 @@ def __init__(
         if app is not None:
             self.init_app(app)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.app = None
+
     @property
     def description(self,):
         return self._description
@@ -128,11 +137,9 @@ def version(self, version: str):
     def init_app(self, app):
         self.app = app
 
-        app.teardown_appcontext(self.teardown)
-
         # Register Flask extension
         app.extensions = getattr(app, "extensions", {})
-        app.extensions[EXTENSION_NAME] = self
+        app.extensions[EXTENSION_NAME] = weakref.ref(self)
 
         # Flask error formatter
         if self.format_flask_exceptions:
@@ -154,12 +161,9 @@ def init_app(self, app):
         self.sockets = Sockets(app)
         self._create_base_sockets()
 
-    def teardown(self, exception):
-        pass
-
     def _create_base_routes(self):
         # Add root representation
-        self.app.add_url_rule(self._complete_url("/", ""), "root", self.root)
+        self.add_view(RootView, self._complete_url("/", ""), endpoint="root")
         # Add thing descriptions
         self.app.register_blueprint(
             docs_blueprint, url_prefix=f"{self.url_prefix}/docs"
@@ -175,19 +179,7 @@ def _create_base_routes(self):
         self.add_view(TaskView, "/tasks/<task_id>", endpoint=TASK_ENDPOINT)
 
     def _create_base_sockets(self):
-        self.sockets.add_view(self._complete_url("", ""), self._socket_handler)
-
-    def _socket_handler(self, ws):
-        # Create a socket subscriber
-        wssub = SocketSubscriber(ws)
-        self.subscribers.add(wssub)
-        logging.info(f"Added subscriber {wssub}")
-        # Start the socket connection handler loop
-        socket_handler_loop(ws)
-        # Remove the subscriber once the loop returns
-        self.subscribers.remove(wssub)
-        logging.info(f"Removed subscriber {wssub}")
-        logging.debug(list(self.subscribers))
+        self.sockets.add_view(self._complete_url("", ""), socket_handler)
 
     # Device stuff
 
@@ -309,6 +301,8 @@ def _register_view(self, app, view, *urls, endpoint=None, **kwargs):
 
         # There might be a better way to do this than _rules_by_endpoint,
         # but I can't find one so this will do for now. Skipping PYL-W0212
+        # FIXME: There is a MASSIVE memory leak or something going on in APISpec!
+        # This is grinding tests to a halt, and is really annoying... Should be fixed.
         flask_rules = app.url_map._rules_by_endpoint.get(endpoint)  # skipcq: PYL-W0212
         for flask_rule in flask_rules:
             self.spec.path(**rule_to_apispec_path(flask_rule, view, self.spec))
@@ -345,8 +339,3 @@ def add_root_link(self, view, rel, kwargs=None, params=None):
         if params is None:
             params = {}
         self.thing_description.add_link(view, rel, kwargs=kwargs, params=params)
-
-    # Description
-    def root(self):
-        """Root representation"""
-        return self.thing_description.to_dict()

labthings/server/logging.py

@@ -1,17 +1,18 @@
+from .find import current_labthing
+
 from logging import StreamHandler
 import datetime
 
 
 class LabThingLogger(StreamHandler):
-    def __init__(self, labthing):
-        StreamHandler.__init__(self)
-        self.labthing = labthing
+    def __init__(self, *args, **kwargs):
+        StreamHandler.__init__(self, *args, **kwargs)
 
     def emit(self, record):
         log_event = self.rest_format_record(record)
 
         # Broadcast to subscribers
-        subscribers = getattr(self.labthing, "subscribers", [])
+        subscribers = getattr(current_labthing(), "subscribers", [])
         for sub in subscribers:
             sub.event_notify(log_event)
 

tests/conftest.py

@@ -190,6 +190,12 @@ def thing(app):
         return thing
 
 
+@pytest.fixture()
+def thing_ctx(thing):
+    with thing.app.app_context():
+        yield thing.app
+
+
 @pytest.fixture()
 def debug_app(request):
 

tests/test_server_find.py

@@ -0,0 +1,75 @@
+from labthings.server import find
+
+from labthings.server.extensions import BaseExtension
+
+
+def test_current_labthing(thing, thing_ctx):
+    with thing_ctx.test_request_context():
+        assert find.current_labthing() is thing
+
+
+def test_current_labthing_explicit_app(thing, thing_ctx):
+    with thing_ctx.test_request_context():
+        assert find.current_labthing(thing.app) is thing
+
+
+def test_current_labthing_missing_app():
+    assert find.current_labthing() is None
+
+
+def test_registered_extensions(thing_ctx):
+    with thing_ctx.test_request_context():
+        assert find.registered_extensions() == {}
+
+
+def test_registered_extensions_explicit_thing(thing):
+    assert find.registered_extensions(thing) == {}
+
+
+def test_registered_components(thing_ctx):
+    with thing_ctx.test_request_context():
+        assert find.registered_components() == {}
+
+
+def test_registered_components_explicit_thing(thing):
+    assert find.registered_components(thing) == {}
+
+
+def test_find_component(thing, thing_ctx):
+    component = type("component", (object,), {})
+    thing.add_component(component, "org.labthings.tests.component")
+
+    with thing_ctx.test_request_context():
+        assert find.find_component("org.labthings.tests.component") == component
+
+
+def test_find_component_explicit_thing(thing):
+    component = type("component", (object,), {})
+    thing.add_component(component, "org.labthings.tests.component")
+
+    assert find.find_component("org.labthings.tests.component", thing) == component
+
+
+def test_find_component_missing_component(thing_ctx):
+    with thing_ctx.test_request_context():
+        assert find.find_component("org.labthings.tests.component") is None
+
+
+def test_find_extension(thing, thing_ctx):
+    extension = BaseExtension("org.labthings.tests.extension")
+    thing.register_extension(extension)
+
+    with thing_ctx.test_request_context():
+        assert find.find_extension("org.labthings.tests.extension") == extension
+
+
+def test_find_extension_explicit_thing(thing):
+    extension = BaseExtension("org.labthings.tests.extension")
+    thing.register_extension(extension)
+
+    assert find.find_extension("org.labthings.tests.extension", thing) == extension
+
+
+def test_find_extension_missing_extesion(thing_ctx):
+    with thing_ctx.test_request_context():
+        assert find.find_extension("org.labthings.tests.extension") is None

tests/test_server_labthing.py

@@ -24,7 +24,9 @@ def test_init_app(app):
     thing = labthing.LabThing()
     thing.init_app(app)
 
-    assert app.extensions.get(EXTENSION_NAME) == thing
+    # Check weakref
+    assert app.extensions.get(EXTENSION_NAME)() == thing
+
     assert app.json_encoder == LabThingsJSONEncoder
     assert 400 in app.error_handler_spec.get(None)
 
@@ -218,11 +220,6 @@ def test_td_add_link_options(thing, view_cls):
     } in thing.thing_description._links
 
 
-def test_root_rep(thing, app_ctx):
-    with app_ctx.test_request_context():
-        assert thing.root() == thing.thing_description.to_dict()
-
-
 def test_description(thing):
     assert thing.description == ""
     thing.description = "description"
@@ -242,10 +239,3 @@ def test_version(thing):
     thing.version = "x.x.x"
     assert thing.version == "x.x.x"
     assert thing.spec.version == "x.x.x"
-
-
-def test_socket_handler(thing, fake_websocket):
-    ws = fake_websocket("", recieve_once=True)
-    thing._socket_handler(ws)
-    # Expect no response
-    assert ws.responses == []