pytests: refactor to use expect_kresd_close

diff --git a/tests/pytests/test_conn_mgmt.py b/tests/pytests/test_conn_mgmt.py
index 6c4a839a3811160eb51f00b827e3d545fb2c838b..6658d9388dda9a808e2c8d6497de9c7734dfcc3a 100644 (file)
--- 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 3b1e33e69c919c49aebfba8f36141b7c9abdd23e..03e0eaf1c40d52e4cffce08b91de5b36035269a3 100644 (file)
--- 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 896840ef72d92ffb8bdd9b847f012e2b157b53b0..97328d454e5d05a51d882450bdf7ccc623ecf401 100644 (file)
--- 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")