From 1897aaf7fa088365b9a597dd93797753fa252f8f Mon Sep 17 00:00:00 2001
From: Guillaume Gay <guillaume@damcb.com>
Date: Fri, 28 May 2021 17:17:10 +0200
Subject: [PATCH 1/4] Assert connection decorator, see #289

---
 src/omero/gateway/__init__.py | 22 +++++++++++++++++++---
 test/unit/test_gateway.py     | 19 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/omero/gateway/__init__.py b/src/omero/gateway/__init__.py
index 0c39a923e..decda0903 100644
--- a/src/omero/gateway/__init__.py
+++ b/src/omero/gateway/__init__.py
@@ -26,7 +26,7 @@
 from past.utils import old_div
 from builtins import object
 import os
-
+from functools import wraps
 import warnings
 from collections import defaultdict
 
@@ -1505,6 +1505,23 @@ def values(self):
         return ()
 
 
+def assertConnected(func):
+    """decorator that raises an exception if  the decorated function is used when the gateway
+    is not connected.
+
+    Assumes the first argument is a BlitzGateway instance, or the function
+    is a BlitzGateway method
+
+    """
+    @wraps(func)
+    def wrapped(conn, *args, **kwargs):
+        if not conn._connected:
+            raise Ice.ConnectionLostException("Trying to call %s while not connected" % func.__name__)
+
+        return func(conn, *args, **kwargs)
+    return wrapped
+
+
 class _BlitzGateway (object):
     """
     Connection wrapper. Handles connecting and keeping the session alive,
@@ -2301,12 +2318,11 @@ def isConnected(self):
 
         :return:    Boolean
         """
-
         return self._connected
 
     ########################
     # # Connection Stuff # #
-
+    @assertConnected
     def getEventContext(self):
         """
         Returns omero_System_ice.EventContext.
diff --git a/test/unit/test_gateway.py b/test/unit/test_gateway.py
index 3f34480ca..37cd23f3b 100644
--- a/test/unit/test_gateway.py
+++ b/test/unit/test_gateway.py
@@ -30,7 +30,7 @@
 import sys
 
 from omero.gateway import BlitzGateway, ImageWrapper, \
-    WellWrapper, LogicalChannelWrapper, OriginalFileWrapper
+    WellWrapper, LogicalChannelWrapper, OriginalFileWrapper, assertConnected
 from omero.model import ImageI, PixelsI, ExperimenterI, EventI, \
     ProjectI, TagAnnotationI, FileAnnotationI, OriginalFileI, \
     MapAnnotationI, NamedValue, PlateI, WellI, \
@@ -40,6 +40,23 @@
 from omero.rtypes import rstring, rtime, rlong, rint, rdouble
 
 
+def test_assertConnected():
+
+    conn = BlitzGateway(
+        host="localhost",
+        username="admin",
+        passwd="omero",
+        secure=True
+    )
+    assert not conn.isConnected()
+    with pytest.raises(Ice.ConnectionLostException):
+        conn.getEventContext()
+
+    assert conn.connect()
+    ctx = conn.getEventContext()
+    assert ctx is not None
+
+
 class MockQueryService(object):
 
     def __init__(self, obj_to_be_returned):

From 535418de2cf72d9dfc94e25717c62a04b1cda791 Mon Sep 17 00:00:00 2001
From: Guillaume Gay <guillaume@damcb.com>
Date: Mon, 31 May 2021 11:59:55 +0200
Subject: [PATCH 2/4] minimal mock gateway to test decorator

---
 test/unit/test_gateway.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/test/unit/test_gateway.py b/test/unit/test_gateway.py
index 37cd23f3b..7eaf89df8 100644
--- a/test/unit/test_gateway.py
+++ b/test/unit/test_gateway.py
@@ -42,19 +42,27 @@
 
 def test_assertConnected():
 
-    conn = BlitzGateway(
-        host="localhost",
-        username="admin",
-        passwd="omero",
-        secure=True
-    )
+    class MockGateway:
+        _connected = False
+
+        def connect(self):
+            self._connected = True
+
+        def isConnected(self):
+            return self._connected
+
+        @assertConnected
+        def foo(self):
+            return 1
+
+    conn = MockGateway()
     assert not conn.isConnected()
     with pytest.raises(Ice.ConnectionLostException):
-        conn.getEventContext()
+        conn.foo()
 
     assert conn.connect()
     ctx = conn.getEventContext()
-    assert ctx is not None
+    assert conn.foo()
 
 
 class MockQueryService(object):

From 2d184dc62efcfcff8a982196923a0ac3ed8c1927 Mon Sep 17 00:00:00 2001
From: Guillaume Gay <guillaume@damcb.com>
Date: Mon, 31 May 2021 14:42:18 +0200
Subject: [PATCH 3/4] return true in mock connect

---
 test/unit/test_gateway.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/unit/test_gateway.py b/test/unit/test_gateway.py
index 7eaf89df8..d86c08e50 100644
--- a/test/unit/test_gateway.py
+++ b/test/unit/test_gateway.py
@@ -47,6 +47,7 @@ class MockGateway:
 
         def connect(self):
             self._connected = True
+            return True
 
         def isConnected(self):
             return self._connected

From 1abfc2cf824481eb5f870501bb6e24dc8dc17786 Mon Sep 17 00:00:00 2001
From: Guillaume Gay <guillaume@damcb.com>
Date: Mon, 31 May 2021 14:46:10 +0200
Subject: [PATCH 4/4] forgot to add before amend

---
 test/unit/test_gateway.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/unit/test_gateway.py b/test/unit/test_gateway.py
index d86c08e50..36e1beb26 100644
--- a/test/unit/test_gateway.py
+++ b/test/unit/test_gateway.py
@@ -62,7 +62,6 @@ def foo(self):
         conn.foo()
 
     assert conn.connect()
-    ctx = conn.getEventContext()
     assert conn.foo()