From 783b0a85266947adb31fd5511cf413b426aa1e60 Mon Sep 17 00:00:00 2001
From: Victor Stinner <victor.stinner@gmail.com>
Date: Mon, 21 Jul 2014 14:05:41 +0200
Subject: [PATCH] Cleanup tests

---
 tests/test_asyncio_trollius.py |  3 +--
 tests/test_tasks.py            | 13 +++++++------
 tests/test_tasks_trollius.py   |  9 +++++----
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/tests/test_asyncio_trollius.py b/tests/test_asyncio_trollius.py
index 4efa0f7..cbcbafc 100644
--- a/tests/test_asyncio_trollius.py
+++ b/tests/test_asyncio_trollius.py
@@ -17,8 +17,7 @@ class TrolliusEventLoopTests(TestCase):
         asyncio.set_event_loop_policy(policy)
         trollius.set_event_loop_policy(policy)
         self.loop = trollius.new_event_loop()
-        asyncio.set_event_loop(self.loop)
-        trollius.set_event_loop(self.loop)
+        policy.set_event_loop(self.loop)
 
     def tearDown(self):
         self.loop.close()
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index bd21710..2c06062 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -4,12 +4,12 @@
 ##
 
 
-import greenio
 import asyncio
-from unittest import TestCase
+import greenio
+import unittest
 
 
-class TaskTests(TestCase):
+class TaskTests(unittest.TestCase):
     def setUp(self):
         asyncio.set_event_loop_policy(greenio.GreenEventLoopPolicy())
         self.loop = asyncio.new_event_loop()
@@ -40,7 +40,7 @@ class TaskTests(TestCase):
         self.assertEqual(fut.result(), 42)
 
     def test_task_yield_from_exception_propagation(self):
-        non_local = {'CHK': 0}
+        CHK = 0
 
         @asyncio.coroutine
         def bar():
@@ -57,10 +57,11 @@ class TaskTests(TestCase):
             try:
                 return (yield from foo())
             except ZeroDivisionError:
-                non_local['CHK'] += 1
+                nonlocal CHK
+                CHK += 1
 
         self.loop.run_until_complete(test())
-        self.assertEqual(non_local['CHK'], 1)
+        self.assertEqual(CHK, 1)
 
     def test_task_yield_from_coroutine(self):
         @asyncio.coroutine
diff --git a/tests/test_tasks_trollius.py b/tests/test_tasks_trollius.py
index bfa3f5d..b440dd8 100644
--- a/tests/test_tasks_trollius.py
+++ b/tests/test_tasks_trollius.py
@@ -16,11 +16,12 @@ except ImportError:
 
 class TrolliusTaskTests(TestCase):
     def setUp(self):
-        trollius.set_event_loop_policy(greenio.GreenTrolliusEventLoopPolicy())
-        self.loop = trollius.new_event_loop()
-        trollius.set_event_loop(self.loop)
+        policy = greenio.GreenTrolliusEventLoopPolicy()
+        trollius.set_event_loop_policy(policy)
         if asyncio is not None:
-            asyncio.set_event_loop(self.loop)
+            asyncio.set_event_loop_policy(policy)
+        self.loop = trollius.new_event_loop()
+        policy.set_event_loop(self.loop)
 
     def tearDown(self):
         self.loop.close()