Make TestFFTThreadSafe tests Py3-compatible

1 files changed, 16 insertions, 7 deletions

diff --git a/numpy/fft/tests/test_fftpack.py b/numpy/fft/tests/test_fftpack.py
index 13db3e682..ba19e5594 100644
--- a/numpy/fft/tests/test_fftpack.py
+++ b/numpy/fft/tests/test_fftpack.py
@@ -1,8 +1,15 @@
 from __future__ import division, absolute_import, print_function
 
 import numpy as np
-from numpy.testing import TestCase, run_module_suite, assert_array_equal, assert_array_almost_equal
-import threading, Queue
+from numpy.testing import TestCase, run_module_suite, assert_array_almost_equal
+from numpy.testing import assert_array_equal
+import threading
+import sys
+if sys.version_info[0] >= 3:
+    import queue
+else:
+    import Queue as queue
+
 
 def fft1(x):
     L = len(x)
@@ -10,6 +17,7 @@ def fft1(x):
     phase = np.arange(L).reshape(-1, 1) * phase
     return np.sum(x*np.exp(phase), axis=1)
 
+
 class TestFFTShift(TestCase):
 
     def test_fft_n(self):
@@ -25,23 +33,24 @@ class TestFFT1D(TestCase):
 
 
 class TestFFTThreadSafe(TestCase):
-    threads = 32
+    threads = 16
     input_shape = (1000, 1000)
 
     def _test_mtsame(self, func, *args):
         def worker(args, q):
             q.put(func(*args))
 
-        q = Queue.Queue()
+        q = queue.Queue()
         expected = func(*args)
 
         # Spin off a bunch of threads to call the same function simultaneously
-        for i in xrange(self.threads):
+        for i in range(self.threads):
             threading.Thread(target=worker, args=(args, q)).start()
 
         # Make sure all threads returned the correct value
-        for i in xrange(self.threads):
-            assert_array_equal(q.get(), expected, 'Function returned wrong value in multithreaded context')
+        for i in range(self.threads):
+            assert_array_equal(q.get(timeout=5), expected,
+                'Function returned wrong value in multithreaded context')
 
     def test_fft(self):
         a = np.ones(self.input_shape) * 1+0j