Ported job calculation to use coroutines for tasks

aiida/backends/tests/__init__.py

@@ -92,6 +92,7 @@
         'work.rmq': ['aiida.backends.tests.work.test_rmq'],
         'work.run': ['aiida.backends.tests.work.run'],
         'work.runners': ['aiida.backends.tests.work.test_runners'],
+        'work.test_transport': ['aiida.backends.tests.work.test_transport'],
         'work.utils': ['aiida.backends.tests.work.utils'],
         'work.work_chain': ['aiida.backends.tests.work.work_chain'],
         'work.workfunctions': ['aiida.backends.tests.work.test_workfunctions'],

aiida/backends/tests/computer.py

@@ -15,8 +15,6 @@
 from aiida.common.exceptions import NotExistent
 
 
-
-
 class TestComputer(AiidaTestCase):
 
     def test_get_transport(self):
@@ -26,17 +24,18 @@ def test_get_transport(self):
         import tempfile
         from aiida.orm import Computer
         from aiida.orm.backend import construct_backend
-        backend = construct_backend()
 
         new_comp = Computer(name='bbb',
-                                hostname='localhost',
-                                transport_type='local',
-                                scheduler_type='direct',
-                                workdir='/tmp/aiida')
+                            hostname='localhost',
+                            transport_type='local',
+                            scheduler_type='direct',
+                            workdir='/tmp/aiida')
         new_comp.store()
 
         # Configure the computer - no parameters for local transport
-        authinfo = backend.authinfos.create(computer=new_comp, user=backend.users.get_automatic_user())
+        authinfo = self.backend.authinfos.create(
+            computer=new_comp,
+            user=self.backend.users.get_automatic_user())
         authinfo.store()
 
         transport = new_comp.get_transport()
@@ -51,20 +50,19 @@ def test_get_transport(self):
     def test_delete(self):
         from aiida.orm import Computer
         new_comp = Computer(name='aaa',
-                                hostname='aaa',
-                                transport_type='local',
-                                scheduler_type='pbspro',
-                                workdir='/tmp/aiida')
+                            hostname='aaa',
+                            transport_type='local',
+                            scheduler_type='pbspro',
+                            workdir='/tmp/aiida')
         new_comp.store()
 
         comp_pk = new_comp.pk
 
         check_computer = Computer.get(comp_pk)
         self.assertEquals(comp_pk, check_computer.pk)
-        
+
         from aiida.orm.computer import delete_computer
         delete_computer(pk=comp_pk)
 
         with self.assertRaises(NotExistent):
             Computer.get(comp_pk)
-

aiida/backends/tests/work/test_transport.py

@@ -0,0 +1,78 @@
+from tornado.gen import coroutine, Return
+
+from aiida.backends.testbase import AiidaTestCase
+from aiida.work.transports import TransportQueue
+
+
+class TestTransportQueue(AiidaTestCase):
+    """ Tests for the transport queue """
+
+    @classmethod
+    def setUpClass(cls, *args, **kwargs):
+        """ Set up a simple authinfo and for later use """
+        super(TestTransportQueue, cls).setUpClass(*args, **kwargs)
+        # Configure the computer - no parameters for local transport
+        # WARNING: This is not deleted as there is no API facing way to do this
+        # it would require a backend specific call
+        cls.authinfo = cls.backend.authinfos.create(
+            computer=cls.computer,
+            user=cls.backend.users.get_automatic_user())
+        cls.authinfo.store()
+
+    def test_simple_request(self):
+        """ Test a simple transport request """
+        queue = TransportQueue()
+        loop = queue.loop()
+
+        @coroutine
+        def test():
+            trans = None
+            with queue.request_transport(self.authinfo) as request:
+                trans = yield request
+                self.assertTrue(trans.is_open)
+            self.assertFalse(trans.is_open)
+
+        loop.run_sync(lambda: test())
+
+    def test_get_transport_nested(self):
+        """ Test nesting calls to get the same transport """
+        transport_queue = TransportQueue()
+        loop = transport_queue.loop()
+
+        @coroutine
+        def nested(queue, authinfo):
+            with queue.request_transport(authinfo) as request1:
+                trans1 = yield request1
+                self.assertTrue(trans1.is_open)
+                with queue.request_transport(authinfo) as request2:
+                    trans2 = yield request2
+                    self.assertIs(trans1, trans2)
+                    self.assertTrue(trans2.is_open)
+
+        loop.run_sync(lambda: nested(transport_queue, self.authinfo))
+
+    def test_get_transport_interleaved(self):
+        """ Test interleaved calls to get the same transport """
+        transport_queue = TransportQueue()
+        loop = transport_queue.loop()
+
+        @coroutine
+        def interleaved(authinfo):
+            with transport_queue.request_transport(authinfo) as trans_future:
+                yield trans_future
+
+        loop.run_sync(lambda: [interleaved(self.authinfo), interleaved(self.authinfo)])
+
+    def test_return_from_context(self):
+        """ Test raising a Return from coroutine context """
+        queue = TransportQueue()
+        loop = queue.loop()
+
+        @coroutine
+        def test():
+            with queue.request_transport(self.authinfo) as request:
+                trans = yield request
+                raise Return(trans.is_open)
+
+        retval = loop.run_sync(lambda: test())
+        self.assertTrue(retval)

aiida/transport/transport.py

@@ -72,7 +72,12 @@ def __enter__(self):
         """
         # Keep track of how many times enter has been called
         if self._enters == 0:
-            self.open()
+            if self.is_open:
+                # Already open, so just add one to the entered counter
+                # this way on the final exit we will not close
+                self._enters += 1
+            else:
+                self.open()
         self._enters += 1
         return self