from operator import add
from time import sleep

from distributed3.core import read, write, rpc, connect
from distributed3.center import Center
from distributed3.worker import Worker


from tornado import gen
from tornado.ioloop import IOLoop


def test_worker():
    @gen.coroutine
    def f():
        c = Center('127.0.0.1', 8007)
        c.listen(c.port)
        a = Worker('127.0.0.1', 8008, c.ip, c.port)
        yield a._start()
        b = Worker('127.0.0.1', 8009, c.ip, c.port)
        yield b._start()

        while len(c.ncores) < 2:
            yield gen.sleep(0.01)

        a_stream = yield connect(a.ip, a.port)
        b_stream = yield connect(b.ip, b.port)

        response = yield rpc(a_stream).compute(key='x', function=add,
                                               args=[1, 2], needed=[],
                                               close=True)
        assert response == b'success'
        assert a.data['x'] == 3
        assert c.who_has['x'] == set([(a.ip, a.port)])

        response = yield rpc(b_stream).compute(key='y', function=add,
                                               args=['x', 10], needed=['x'])
        assert response == b'success'
        assert b.data['y'] == 13
        assert c.who_has['y'] == set([(b.ip, b.port)])

        def bad_func():
            1 / 0

        response = yield rpc(b_stream).compute(key='z', function=bad_func,
                                               args=(), needed=(), close=True)
        assert response == b'error'
        assert isinstance(b.data['z'], ZeroDivisionError)

        a_stream.close()
        yield a._close()

        assert a.address not in c.ncores and b.address in c.ncores

        assert list(c.ncores.keys()) == [(b.ip, b.port)]

        b_stream.close()
        yield b._close()

    IOLoop.current().run_sync(f)

"""

def test_workers_update_center():
    c = Center('127.0.0.1', 8007, loop=loop)
    a = Worker('127.0.0.1', 8008, c.ip, c.port, loop=loop)

    @asyncio.coroutine
    def f():
        a_reader, a_writer = yield from connect(a.ip, a.port, loop=loop)

        yield from rpc(a_reader, a_writer).update_data(data={'x': 1, 'y': 2})

        assert a.data == {'x': 1, 'y': 2}
        assert c.who_has == {'x': {(a.ip, a.port)},
                             'y': {(a.ip, a.port)}}
        assert c.has_what == {(a.ip, a.port): {'x', 'y'}}

        yield from rpc(a_reader, a_writer).delete_data(keys=['x'], close=True)

        yield from a._close()
        yield from c._close()

    loop.run_until_complete(asyncio.gather(c.go(), a.go(), f(), loop=loop))


def test_close():
    c = Center('127.0.0.1', 8007, loop=loop)
    a = Worker('127.0.0.1', 8008, c.ip, c.port, loop=loop)

    @asyncio.coroutine
    def f():
        while len(c.ncores) < 1:
            yield from asyncio.sleep(0.01, loop=loop)

        assert a.status == 'running'
        yield from rpc(a.ip, a.port, loop=loop).terminate()
        assert a.status == 'closed'

        assert c.status == 'running'
        yield from rpc(c.ip, c.port, loop=loop).terminate()
        assert c.status == 'closed'

    loop.run_until_complete(asyncio.gather(c.go(), a.go(), f(), loop=loop))
"""

@@ -98,16 +98,56 @@
 connect%0A
+from distributed3.utils import ignoring%0A
 from dis
@@ -286,20 +286,23 @@
 def 
+_
 test_
-work
+clust
 er(
+f
 ):%0A 
@@ -319,33 +319,33 @@
 routine%0A    def 
-f
+g
 ():%0A        c = 
@@ -366,17 +366,17 @@
 0.1', 80
-0
+1
 7)%0A     
@@ -421,33 +421,33 @@
 ('127.0.0.1', 80
-0
+1
 8, c.ip, c.port)
@@ -437,32 +437,42 @@
 18, c.ip, c.port
+, ncores=1
 )%0A        yield 
@@ -516,17 +516,17 @@
 0.1', 80
-0
+1
 9, c.ip,
@@ -528,24 +528,34 @@
 c.ip, c.port
+, ncores=1
 )%0A        yi
@@ -638,16 +638,314 @@
 (0.01)%0A%0A
+        try:%0A            yield f(c, a, b)%0A        finally:%0A            with ignoring():%0A                yield a._close()%0A            with ignoring():%0A                yield b._close()%0A            c.stop()%0A%0A    IOLoop.current().run_sync(g)%0A%0A%0Adef test_worker():%0A    @gen.coroutine%0A    def f(c, a, b):%0A
         
@@ -2159,38 +2159,28 @@
     
-IOLoop.current().run_sync
+_test_cluster
 (f)%0A%0A
+%0A%0A
 %22%22%22%0A

refactor worker test to use test_cluster