pytests: import test_query_flood_close

diff --git a/tests/pytests/conftest.py b/tests/pytests/conftest.py
index 67ef5179d6a5d06b98a0097da933a076e3a09fb8..cbb89a4894cfa46ec3cd75776ea79bb9f94e69ee 100644 (file)
--- a/tests/pytests/conftest.py
+++ b/tests/pytests/conftest.py
@@ -54,5 +54,15 @@ def kresd(tmpdir):
     'ip_tls_socket',
     'ip6_tls_socket',
 ])
-def kresd_sock(request, kresd):
-    return getattr(kresd, request.param)()
+def make_kresd_sock(request, kresd):
+    sock_func = getattr(kresd, request.param)
+
+    def _make_kresd_sock():
+        return sock_func()
+
+    return _make_kresd_sock
+
+
+@pytest.fixture
+def kresd_sock(make_kresd_sock):
+    return make_kresd_sock()

diff --git a/tests/pytests/test_conn_mgmt.py b/tests/pytests/test_conn_mgmt.py
index 4ef9080951d7f83d9001b38f01da8db1f06d9e60..ca796cb11ac96a9bc7e00ccd2a85d275a493725c 100644 (file)
--- a/tests/pytests/test_conn_mgmt.py
+++ b/tests/pytests/test_conn_mgmt.py
@@ -1,5 +1,6 @@
 """TCP Connection Management tests"""
 
+import struct
 import time
 
 import pytest
@@ -117,3 +118,30 @@ def test_ignore_jumbo_message(kresd_sock):
 
     answer = utils.receive_parse_answer(kresd_sock)
     assert answer.id == msgid2
+
+
+def test_query_flood_close(make_kresd_sock):
+    """
+    Test floods resolver with queries and closes the connection.
+
+    Expected: resolver must not crash
+    """
+    def flood_buffer(msgcount):
+        flood_buff = bytes()
+        msgbuff, _ = utils.get_msgbuff()
+        noid_msgbuff = msgbuff[2:]
+
+        def gen_msg(msgid):
+            return struct.pack("!H", len(msgbuff)) + struct.pack("!H", msgid) + noid_msgbuff
+
+        for i in range(msgcount):
+            flood_buff += gen_msg(i)
+        return flood_buff
+
+    buff = flood_buffer(10000)
+    sock1 = make_kresd_sock()
+    sock1.sendall(buff)
+    sock1.close()
+
+    sock2 = make_kresd_sock()
+    utils.ping_alive(sock2)