1 files changed, 24 insertions, 29 deletions
diff --git a/gr-blocks/python/qa_argmax.py b/gr-blocks/python/qa_argmax.py
index 1fbda7a1f9..ec82b71cd4 100644
--- a/gr-blocks/python/qa_argmax.py
+++ b/gr-blocks/python/qa_argmax.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 #
-# Copyright 2007,2010 Free Software Foundation, Inc.
+# Copyright 2007,2010,2013 Free Software Foundation, Inc.
 #
 # This file is part of GNU Radio
 #
@@ -24,17 +24,14 @@ from gnuradio import gr, gr_unittest
 import blocks_swig as blocks
 import math
 
+class test_arg_max(gr_unittest.TestCase):
 
-class test_arg_max (gr_unittest.TestCase):
+    def setUp(self):
+        self.tb = gr.top_block()
 
-    def setUp (self):
-        self.tb = gr.top_block ()
-
-
-    def tearDown (self):
+    def tearDown(self):
         self.tb = None
 
-
     def test_001(self):
         tb = self.tb
 
@@ -42,36 +39,34 @@ class test_arg_max (gr_unittest.TestCase):
         src2_data = (0,0.0,3.0,0,10,0)
         src3_data = (0,0.0,3.0,0,1,0)
 
-        src1 = gr.vector_source_f (src1_data)
+        src1 = gr.vector_source_f(src1_data)
         s2v1 = blocks.stream_to_vector(gr.sizeof_float, len(src1_data))
-        tb.connect( src1, s2v1 )
+        tb.connect(src1, s2v1)
 
-        src2 = gr.vector_source_f (src2_data)
+        src2 = gr.vector_source_f(src2_data)
         s2v2 = blocks.stream_to_vector(gr.sizeof_float, len(src1_data))
-        tb.connect( src2, s2v2 )
+        tb.connect(src2, s2v2)
 
-        src3 = gr.vector_source_f (src3_data)
+        src3 = gr.vector_source_f(src3_data)
         s2v3 = blocks.stream_to_vector(gr.sizeof_float, len(src1_data))
-        tb.connect( src3, s2v3 )
-
-        dst1 = gr.vector_sink_s ()
-        dst2 = gr.vector_sink_s ()
-        argmax = gr.argmax_fs (len(src1_data))
-
-        tb.connect (s2v1, (argmax, 0))
-        tb.connect (s2v2, (argmax, 1))
-        tb.connect (s2v3, (argmax, 2))
+        tb.connect(src3, s2v3)
 
-        tb.connect ((argmax,0), dst1)
-        tb.connect ((argmax,1), dst2)
+        dst1 = gr.vector_sink_s()
+        dst2 = gr.vector_sink_s()
+        argmax = blocks.argmax_fs(len(src1_data))
 
-        tb.run ()
-        index = dst1.data ()
-        source = dst2.data ()
-        self.assertEqual ( index, (4,))
-        self.assertEqual ( source, (0,))
+        tb.connect(s2v1, (argmax, 0))
+        tb.connect(s2v2, (argmax, 1))
+        tb.connect(s2v3, (argmax, 2))
 
+        tb.connect((argmax,0), dst1)
+        tb.connect((argmax,1), dst2)
 
+        tb.run()
+        index = dst1.data()
+        source = dst2.data()
+        self.assertEqual(index, (4,))
+        self.assertEqual(source, (0,))
 
 if __name__ == '__main__':
     gr_unittest.run(test_arg_max, "test_arg_max.xml")