""" _utest.py

Unit tests for yoton.

"""

import sys, os, time
import unittest

# Go up one directory and then import the yoton package
os.chdir("..")
os.chdir("..")
sys.path.insert(0, ".")

# Import yoton from there
import yoton
from yoton.misc import long


class Tester(unittest.TestCase):
    """Perform a set of tests using three contexts and
    each of the six channels at each context.

    """

    def setUp(self):
        # Create contexts
        self._context1 = c1 = yoton.Context()
        self._context2 = c2 = yoton.Context()
        self._context3 = c3 = yoton.Context()

        # Create pub sub channels
        self._channel_pub1 = yoton.PubChannel(c1, "foo")
        self._channel_pub2 = yoton.PubChannel(c2, "foo")
        self._channel_pub3 = yoton.PubChannel(c3, "foo")
        #
        self._channel_sub1 = yoton.SubChannel(c1, "foo")
        self._channel_sub2 = yoton.SubChannel(c2, "foo")
        self._channel_sub3 = yoton.SubChannel(c3, "foo")

        # Create req rep channels
        self._channel_req1 = yoton.ReqChannel(c1, "bar")
        self._channel_req2 = yoton.ReqChannel(c2, "bar")
        self._channel_req3 = yoton.ReqChannel(c3, "bar")
        #
        self._channel_rep1 = yoton.RepChannel(c1, "bar")
        self._channel_rep2 = yoton.RepChannel(c2, "bar")
        self._channel_rep3 = yoton.RepChannel(c3, "bar")

        # Create state channels
        self._channel_state1 = yoton.StateChannel(c1, "spam")
        self._channel_state2 = yoton.StateChannel(c2, "spam")
        self._channel_state3 = yoton.StateChannel(c3, "spam")

    def tearDown(self):
        self._context1.close()
        self._context2.close()
        self._context3.close()

    def test_connecting(self):
        # Connect
        self._context1.bind("localhost:test1")
        self._context2.connect("localhost:test1")
        time.sleep(0.1)  # Give time for binding to finish

        # Test if connected
        self.assertEqual(self._context1.connection_count, 1)
        self.assertEqual(self._context2.connection_count, 1)
        self.assertEqual(self._context3.connection_count, 0)

        # Connect more
        self._context1.bind("localhost:test2")
        self._context3.connect("localhost:test2")
        time.sleep(0.1)  # Give time for binding to finish

        # Test if connected
        self.assertEqual(self._context1.connection_count, 2)
        self.assertEqual(self._context2.connection_count, 1)
        self.assertEqual(self._context3.connection_count, 1)

    def test_closing_channel(self):
        # Connect
        self._context1.bind("localhost:test")
        self._context2.connect("localhost:test")

        # Send data
        self._channel_pub1.send("hello")

        # Receive data
        self.assertEqual(self._channel_sub2.recv(), "hello")
        self.assertEqual(self._channel_sub2.recv(False), "")

        # Close channel
        self._channel_sub2.close()
        self.assertEqual(self._channel_sub2.recv(), "")
        self.assertTrue(self._channel_sub2.closed)

    def test_pub_sub(self):
        # also test hopping

        # Connect
        self._context1.bind("localhost:test1")
        self._context2.connect("localhost:test1")
        self._context2.bind("localhost:test2")
        self._context3.connect("localhost:test2")

        # Send messages
        self._channel_pub1.send("msg1")
        time.sleep(0.1)
        self._channel_pub2.send("msg2")
        time.sleep(0.1)

        # Receive msssages
        self.assertEqual(self._channel_sub2.recv(), "msg1")
        self.assertEqual(self._channel_sub3.recv(), "msg1")
        self.assertEqual(self._channel_sub3.recv(), "msg2")

    def test_pub_sub_select_channel(self):
        # Connect
        self._context1.bind("localhost:test1")
        self._context2.connect("localhost:test1")

        # Create channels
        pub1 = yoton.PubChannel(self._context1, "foo1")
        pub2 = yoton.PubChannel(self._context1, "foo2")
        #
        sub1 = yoton.SubChannel(self._context2, "foo1")
        sub2 = yoton.SubChannel(self._context2, "foo2")

        # Send a bunch of messages
        ii = [0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 1, 0, 1]
        for i in range(len(ii)):
            pub = [pub1, pub2][ii[i]]
            pub.send(str(i))

        time.sleep(0.1)

        # Receive in right order
        count = 0
        while True:
            sub = yoton.select_sub_channel(sub1, sub2)
            if sub:
                i = int(sub.recv())
                self.assertEqual(i, count)
                count += 1
            else:
                break

        # Test count
        self.assertEqual(count, len(ii))

    def test_pubstate_substate(self):
        # also test hopping

        # Set status before connecting
        self._channel_state1.send("status1")

        # Connect
        self._context1.bind("localhost:test1")
        self._context2.connect("localhost:test1")
        self._context2.bind("localhost:test2")
        self._context3.connect("localhost:test2")
        time.sleep(0.1)

        # Receive msssages
        self.assertEqual(self._channel_state2.recv(), "status1")
        self.assertEqual(self._channel_state2.recv(), "status1")
        self.assertEqual(self._channel_state3.recv(), "status1")
        self.assertEqual(self._channel_state3.recv(), "status1")

        # Update status
        self._channel_state1.send("status2")
        time.sleep(0.1)

        # Receive msssages (recv multiple times should return last status)
        self.assertEqual(self._channel_state2.recv(), "status2")
        self.assertEqual(self._channel_state2.recv(), "status2")
        self.assertEqual(self._channel_state3.recv(), "status2")
        self.assertEqual(self._channel_state3.recv(), "status2")

    def test_req_rep1(self):
        # Connect
        self._context1.bind("localhost:test1")
        self._context2.connect("localhost:test1")
        self._context2.bind("localhost:test2")
        self._context3.connect("localhost:test2")

        # Turn on one replier
        self._channel_rep3.set_mode("thread")

        # Do requests
        foo, id1 = self._channel_req1.echo("foo").result(1)
        bar, id2 = self._channel_req1.echo("bar").result(1)

        # Check reply
        self.assertEqual(foo, "foo")
        self.assertEqual(bar, "bar")
        self.assertEqual(id1, hex(self._context3.id))
        self.assertEqual(id2, hex(self._context3.id))

    def test_req_rep2(self):
        # Test multiple requesters and multiple repliers

        # Connect
        if True:
            self._context1.bind("localhost:test1")
            self._context3.connect("localhost:test1")
            #
            self._context1.bind("localhost:test2")
            self._context2.connect("localhost:test2")
        else:
            freeNode = yoton.Context()
            #
            self._context1.bind("localhost:test1")
            freeNode.connect("localhost:test1")
            freeNode.bind("localhost:test1f")
            self._context3.connect("localhost:test1f")
            #
            self._context1.bind("localhost:test2")
            self._context2.connect("localhost:test2")
        time.sleep(0.1)

        # Turn on 2 requesters and 3 repliers
        self._channel_rep1.set_mode("thread")  # threaded because simulating
        self._channel_rep2.set_mode("thread")  # different processes
        self._channel_rep3.set_mode("thread")

        # Define and register reply handler
        def reply_handler(future):
            reply = future.result()
            reqnr = future.reqnr

            echo, id = reply
            if echo.lower() == "stop":
                yoton.stop_event_loop()
            else:
                self.assertEqual(echo[:3], "msg")
                contextnr = {
                    self._context1.id: 1,
                    self._context2.id: 2,
                    self._context3.id: 3,
                }[long(id, 16)]
                print(
                    "request %s from %i handled by context %i."
                    % (echo, reqnr, contextnr)
                )

        # Get echo functions
        echoFun1 = self._channel_req1.echo
        echoFun2 = self._channel_req2.echo

        # Send requests on req 1
        sleepTimes = [1, 0.1, 1, 0.6, 0.6, 0.6, 0.6, 0.6]
        for i in range(len(sleepTimes)):
            f = echoFun1("msg%i" % i, sleepTimes[i])
            f.add_done_callback(reply_handler)
            f.reqnr = 1

        # Send requests on req 2
        sleepTimes = [0.4, 0.4, 0.4, 0.4, 0.4]
        for i in range(len(sleepTimes)):
            f = echoFun2("msg%i" % i, sleepTimes[i])
            f.add_done_callback(reply_handler)
            f.reqnr = 2

        # Stop
        f = echoFun1("stop")
        f.add_done_callback(reply_handler)
        f.reqnr = 1

        # Enter event loop
        yoton.start_event_loop()


if __name__ == "__main__":
    suite = unittest.TestLoader().loadTestsFromTestCase(Tester)
    unittest.TextTestRunner(verbosity=2).run(suite)
    print("")
