From 556c511fc348bbed123792049f81df757afdf3f6 Mon Sep 17 00:00:00 2001
From: Andrej Rode <mail@andrejro.de>
Date: Tue, 4 Dec 2018 15:23:39 +0100
Subject: [PATCH 5/5] blocks: fix the udp_source_sink test flakiness

Provide one unit test for the source and sink with the raw python
socket module.
Fix the tests on IPv4 only to avoid mismatched connections on IPv4 &
IPv6.

Fixes #2210.
---
 gr-blocks/python/blocks/qa_udp_source_sink.py | 72 +++++++++++++++++++++------
 1 file changed, 58 insertions(+), 14 deletions(-)

diff --git a/gr-blocks/python/blocks/qa_udp_source_sink.py b/gr-blocks/python/blocks/qa_udp_source_sink.py
index 905b9e08a0..aba09b9f62 100644
--- a/gr-blocks/python/blocks/qa_udp_source_sink.py
+++ b/gr-blocks/python/blocks/qa_udp_source_sink.py
@@ -21,9 +21,22 @@
 #
 
 from gnuradio import gr, gr_unittest, blocks
+import numpy
 import os
+import socket
+import time
+
+from threading import Timer, Thread
+
+
+def recv_data(sock, result):
+    while True:
+        data = sock.recv(4*1000)
+        if len(data) == 0:
+            break
+        real_data = numpy.frombuffer(data, dtype=numpy.float32)
+        result.extend(list(real_data))
 
-from threading import Timer
 
 class test_udp_sink_source(gr_unittest.TestCase):
 
@@ -51,36 +64,67 @@ class test_udp_sink_source(gr_unittest.TestCase):
         self.tb_snd.run()
         udp_snd.disconnect()
 
+
         udp_snd.connect('localhost', port+1)
         src.rewind()
         self.tb_snd.run()
 
-    def test_002(self):
+
+    def test_sink_001(self):
         port = 65520
 
         n_data = 100
         src_data = [float(x) for x in range(n_data)]
         expected_result = tuple(src_data)
+
+        recvsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        recvsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        recvsock.bind(('127.0.0.1', port))
+
+        result = []
+        t = Thread(target=recv_data, args=(recvsock, result))
+        t.start()
+
         src = blocks.vector_source_f(src_data, False)
-        udp_snd = blocks.udp_sink(gr.sizeof_float, 'localhost', port)
+        udp_snd = blocks.udp_sink(gr.sizeof_float, '127.0.0.1', port)
         self.tb_snd.connect(src, udp_snd)
 
-        udp_rcv = blocks.udp_source(gr.sizeof_float, 'localhost', port)
-        dst = blocks.vector_sink_f()
-        self.tb_rcv.connect(udp_rcv, dst)
 
         self.tb_rcv.start()
         self.tb_snd.run()
         udp_snd.disconnect()
-        self.timeout = False
-        q = Timer(2.0,self.stop_rcv)
-        q.start()
+        t.join()
+        recvsock.close()
+
+        self.assertEqual(expected_result, tuple(result))
+
+    def test_source_001(self):
+        port = 65520
+
+        n_data = 100
+        src_data = [float(x) for x in range(n_data)]
+        expected_result = tuple(src_data)
+        send_data = numpy.array(src_data, dtype=numpy.float32)
+        send_data = send_data.tobytes()
+
+        udp_rcv = blocks.udp_source(gr.sizeof_float, '127.0.0.1', port)
+        dst = blocks.vector_sink_f()
+        self.tb_rcv.connect(udp_rcv, dst)
+        self.tb_rcv.start()
+        time.sleep(1.0)
+        sendsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sendsock.sendto(send_data, ('127.0.0.1', port))
+        time.sleep(1.0)
+        sendsock.sendto(b'', ('127.0.0.1', port))
+        sendsock.sendto(b'', ('127.0.0.1', port))
+        sendsock.sendto(b'', ('127.0.0.1', port))
         self.tb_rcv.wait()
-        q.cancel()
+        sendsock.close()
+        recv_data = tuple(dst.data())
+
+        self.assertEqual(expected_result, recv_data)
+
 
-        result_data = dst.data()
-        self.assertEqual(expected_result, result_data)
-        self.assert_(not self.timeout)
 
     def test_003(self):
         port = 65530
@@ -111,7 +155,7 @@ class test_udp_sink_source(gr_unittest.TestCase):
 
         result_data = dst.data()
         self.assertEqual(expected_result, result_data)
-        self.assert_(self.timeout)  # source ignores EOF?
+        self.assertTrue(self.timeout)  # source ignores EOF?
 
     def stop_rcv(self):
         self.timeout = True
-- 
2.11.0

