From f2dcc50515d771bab5eaba7ba0fe8999a195eccf Mon Sep 17 00:00:00 2001
From: Victor Stinner <victor.stinner@gmail.com>
Date: Mon, 21 Jul 2014 14:22:23 +0200
Subject: [PATCH] Run test_socket on asyncio and trollius

---
 greenio/__init__.py          |  3 ++-
 greenio/socket.py            |  6 ++---
 tests/test_socket.py         | 46 +++++++++++++++++++++++++++++-------
 tests/test_tasks.py          |  7 +++---
 tests/test_tasks_trollius.py |  2 +-
 5 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/greenio/__init__.py b/greenio/__init__.py
index 9189449..393d10c 100644
--- a/greenio/__init__.py
+++ b/greenio/__init__.py
@@ -150,7 +150,8 @@ if trollius is not None:
         class GreenTrolliusTask(_GreenTaskMixin, trollius.Task):
             pass
 
-        class GreenTrolliusUnixSelectorLoop(_GreenLoopMixin, trollius.SelectorEventLoop):
+        class GreenTrolliusUnixSelectorLoop(_GreenLoopMixin,
+                                            trollius.SelectorEventLoop):
             def create_task(self, coro):
                 return GreenTrolliusTask(coro)
 
diff --git a/greenio/socket.py b/greenio/socket.py
index 8451a2b..9a8dbeb 100644
--- a/greenio/socket.py
+++ b/greenio/socket.py
@@ -14,7 +14,7 @@ from socket import error, SOCK_STREAM
 from socket import socket as std_socket
 
 from . import yield_from
-from . import GreenUnixSelectorLoop
+from . import _GreenLoopMixin
 
 
 class socket:
@@ -30,8 +30,8 @@ class socket:
         try:
             self._sock.setblocking(False)
             self._loop = asyncio.get_event_loop()
-            assert isinstance(self._loop, GreenUnixSelectorLoop), \
-                'GreenUnixSelectorLoop event loop is required'
+            assert isinstance(self._loop, _GreenLoopMixin), \
+                'greenio event loop is required'
         except:
             if own_sock is not None:
                 # An unexpected error has occurred. Close the
diff --git a/tests/test_socket.py b/tests/test_socket.py
index edf9874..fd16215 100644
--- a/tests/test_socket.py
+++ b/tests/test_socket.py
@@ -6,10 +6,19 @@
 
 try:
     import asyncio
-    from unittest import TestCase
 except ImportError:
-    import trollius as asyncio
+    asyncio = None
+try:
+    import trollius
+except ImportError:
+    trollius = None
+if asyncio is None and trollius is None:
+    raise ImportError("asyncio and trollius modules are missing")
+
+try:
     from trollius.test_utils import TestCase
+except ImportError:
+    from unittest import TestCase
 
 import greenio
 import greenio.socket as greensocket
@@ -17,21 +26,24 @@ import greenio.socket as greensocket
 import socket as std_socket
 
 
-class SocketTests(TestCase):
+class SocketMixin(object):
+    asyncio = None
+    event_loop_policy = greenio.GreenEventLoopPolicy
 
     def setUp(self):
-        asyncio.set_event_loop_policy(greenio.GreenEventLoopPolicy())
-        self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self.loop)
+        policy = self.event_loop_policy()
+        self.asyncio.set_event_loop_policy(policy)
+        self.loop = policy.new_event_loop()
+        policy.set_event_loop(self.loop)
 
     def tearDown(self):
         self.loop.close()
-        asyncio.set_event_loop_policy(None)
+        self.asyncio.set_event_loop_policy(None)
 
     def test_socket_wrong_event_loop(self):
-        loop = asyncio.DefaultEventLoopPolicy().new_event_loop()
+        loop = self.asyncio.DefaultEventLoopPolicy().new_event_loop()
         self.addCleanup(loop.close)
-        asyncio.set_event_loop(loop)
+        self.asyncio.set_event_loop(loop)
         self.assertRaises(AssertionError, greensocket.socket)
 
     def test_socket_docs(self):
@@ -187,3 +199,19 @@ class SocketTests(TestCase):
             greenio.task(server)(greensocket.socket))
         thread.join(1)
         self.assertEqual(non_local['check'], 1)
+
+if asyncio is not None:
+    class SocketTests(SocketMixin, TestCase):
+        asyncio = asyncio
+        event_loop_policy = greenio.GreenEventLoopPolicy
+
+if trollius is not None:
+    class TrolliusSocketTests(SocketMixin, TestCase):
+        asyncio = trollius
+        event_loop_policy = greenio.GreenTrolliusEventLoopPolicy
+
+        def setUp(self):
+            super(TrolliusSocketTests, self).setUp()
+            if asyncio is not None:
+                policy = trollius.get_event_loop_policy()
+                asyncio.set_event_loop_policy(policy)
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index 2c06062..426d1fb 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -11,9 +11,10 @@ import unittest
 
 class TaskTests(unittest.TestCase):
     def setUp(self):
-        asyncio.set_event_loop_policy(greenio.GreenEventLoopPolicy())
-        self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self.loop)
+        policy = greenio.GreenEventLoopPolicy()
+        asyncio.set_event_loop_policy(policy)
+        self.loop = policy.new_event_loop()
+        policy.set_event_loop(self.loop)
 
     def tearDown(self):
         self.loop.close()
diff --git a/tests/test_tasks_trollius.py b/tests/test_tasks_trollius.py
index b440dd8..6eece0b 100644
--- a/tests/test_tasks_trollius.py
+++ b/tests/test_tasks_trollius.py
@@ -20,7 +20,7 @@ class TrolliusTaskTests(TestCase):
         trollius.set_event_loop_policy(policy)
         if asyncio is not None:
             asyncio.set_event_loop_policy(policy)
-        self.loop = trollius.new_event_loop()
+        self.loop = policy.new_event_loop()
         policy.set_event_loop(self.loop)
 
     def tearDown(self):