small cleanup of test_nanny

This wasn't actually important, but was part of a small experiment that
is a bit cleaner than what we had before.

distributed/core.py

@@ -697,6 +697,12 @@ def __enter__(self):
     def __exit__(self, *args):
         asyncio.ensure_future(self.close_rpc())
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        await self.close_rpc()
+
     def __del__(self):
         if self.status != "closed":
             rpc.active.discard(self)

distributed/tests/test_nanny.py

@@ -30,33 +30,30 @@
 
 
 @gen_cluster(nthreads=[])
-def test_nanny(s):
-    n = yield Nanny(s.address, nthreads=2, loop=s.loop)
-
-    with rpc(n.address) as nn:
-        assert n.is_alive()
-        assert s.nthreads[n.worker_address] == 2
-        assert s.workers[n.worker_address].nanny == n.address
-
-        yield nn.kill()
-        assert not n.is_alive()
-        assert n.worker_address not in s.nthreads
-        assert n.worker_address not in s.workers
-
-        yield nn.kill()
-        assert not n.is_alive()
-        assert n.worker_address not in s.nthreads
-        assert n.worker_address not in s.workers
-
-        yield nn.instantiate()
-        assert n.is_alive()
-        assert s.nthreads[n.worker_address] == 2
-        assert s.workers[n.worker_address].nanny == n.address
-
-        yield nn.terminate()
-        assert not n.is_alive()
-
-    yield n.close()
+async def test_nanny(s):
+    async with Nanny(s.address, nthreads=2, loop=s.loop) as n:
+        async with rpc(n.address) as nn:
+            assert n.is_alive()
+            assert s.nthreads[n.worker_address] == 2
+            assert s.workers[n.worker_address].nanny == n.address
+
+            await nn.kill()
+            assert not n.is_alive()
+            assert n.worker_address not in s.nthreads
+            assert n.worker_address not in s.workers
+
+            await nn.kill()
+            assert not n.is_alive()
+            assert n.worker_address not in s.nthreads
+            assert n.worker_address not in s.workers
+
+            await nn.instantiate()
+            assert n.is_alive()
+            assert s.nthreads[n.worker_address] == 2
+            assert s.workers[n.worker_address].nanny == n.address
+
+            await nn.terminate()
+            assert not n.is_alive()
 
 
 @gen_cluster(nthreads=[])