From: Tomas Krizek <tomas.krizek@nic.cz>
Date: Wed, 14 Nov 2018 13:43:48 +0000 (+0100)
Subject: pytests: refactor to use expect_kresd_close
X-Git-Tag: v3.2.0~18^2~42
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=cee15631a07f5aa0ef4684bc37ee3a49b8dd513f;p=thirdparty%2Fknot-resolver.git

pytests: refactor to use expect_kresd_close
---

diff --git a/tests/pytests/test_conn_mgmt.py b/tests/pytests/test_conn_mgmt.py
index 6c4a839a3..6658d9388 100644
--- a/tests/pytests/test_conn_mgmt.py
+++ b/tests/pytests/test_conn_mgmt.py
@@ -2,8 +2,6 @@
 
 import time
 
-import pytest
-
 import utils
 
 
@@ -64,11 +62,8 @@ def test_close(kresd_sock):
     """
     time.sleep(utils.MAX_TIMEOUT)
 
-    with pytest.raises(BrokenPipeError, message="kresd didn't close the connection"):
-        try:
-            utils.ping_alive(kresd_sock)
-        except ConnectionResetError:
-            pytest.skip('TCP RST')
+    with utils.expect_kresd_close():
+        utils.ping_alive(kresd_sock)
 
 
 def test_slow_lorris_attack(kresd_sock):
@@ -79,11 +74,8 @@ def test_slow_lorris_attack(kresd_sock):
     """
     buff, _ = utils.get_msgbuff()
 
-    with pytest.raises(BrokenPipeError, message="kresd didn't close the connection"):
-        try:
-            for i in range(len(buff)):
-                b = buff[i:i+1]
-                kresd_sock.send(b)
-                time.sleep(1)
-        except ConnectionResetError:
-            pytest.skip('TCP RST')
+    with utils.expect_kresd_close():
+        for i in range(len(buff)):
+            b = buff[i:i+1]
+            kresd_sock.send(b)
+            time.sleep(1)
diff --git a/tests/pytests/test_tcp_prefix.py b/tests/pytests/test_tcp_prefix.py
index 3b1e33e69..03e0eaf1c 100644
--- a/tests/pytests/test_tcp_prefix.py
+++ b/tests/pytests/test_tcp_prefix.py
@@ -16,8 +16,6 @@ TCP_DEFER_ACCEPT, ...), kresd should close the connection.
 
 import time
 
-import pytest
-
 import utils
 
 
@@ -31,12 +29,9 @@ def send_incorrect_repeatedly(sock, buff, delay=1):
     """
     end_time = time.time() + utils.MAX_TIMEOUT
 
-    with pytest.raises(BrokenPipeError, message="kresd didn't close connection"):
+    with utils.expect_kresd_close():
         while time.time() < end_time:
-            try:
-                sock.sendall(buff)
-            except ConnectionResetError:
-                pytest.skip("kresd closed connection with TCP RST")
+            sock.sendall(buff)
             time.sleep(delay)
 
 
diff --git a/tests/pytests/utils.py b/tests/pytests/utils.py
index 896840ef7..97328d454 100644
--- a/tests/pytests/utils.py
+++ b/tests/pytests/utils.py
@@ -1,8 +1,10 @@
+from contextlib import contextmanager
 import struct
 import random
 
 import dns
 import dns.message
+import pytest
 
 
 # default net.tcp_in_idle is 10s, TCP_DEFER_ACCEPT 3s, some extra for
@@ -79,3 +81,12 @@ def ping_alive(sock):
     sock.sendall(buff)
     answer = receive_parse_answer(sock)
     return answer.id == msgid
+
+
+@contextmanager
+def expect_kresd_close():
+    with pytest.raises(BrokenPipeError, message="kresd didn't close the connection"):
+        try:
+            yield
+        except ConnectionResetError:
+            pytest.skip("kresd closed connection with TCP RST")