From f55e1a18723a3b3aca34e9a4b8b878860f951e52 Mon Sep 17 00:00:00 2001
From: Eugene Crosser <crosser@average.org>
Date: Thu, 18 Aug 2022 11:56:38 +0200
Subject: [PATCH] tests: separate fuzz tests in two modules

---
 test/common.py           | 28 ++++++++++++++++++++++-
 test/test_fuzz.py        | 48 ----------------------------------------
 test/test_fuzz_stream.py | 18 +++++++++++++++
 test/test_fuzz_zx303.py  | 18 +++++++++++++++
 4 files changed, 63 insertions(+), 49 deletions(-)
 delete mode 100644 test/test_fuzz.py
 create mode 100644 test/test_fuzz_stream.py
 create mode 100644 test/test_fuzz_zx303.py

diff --git a/test/common.py b/test/common.py
index 7a86310..131169c 100644
--- a/test/common.py
+++ b/test/common.py
@@ -9,18 +9,22 @@ from multiprocessing import Process
 from os import kill, unlink
 from signal import SIGINT
 from socket import (
+    AF_INET,
     AF_INET6,
+    getaddrinfo,
     MSG_DONTWAIT,
     SOCK_DGRAM,
+    SOCK_STREAM,
     SOL_SOCKET,
     SO_REUSEADDR,
     socket,
     SocketType,
 )
 from sys import exit, stderr
+from random import Random
 from tempfile import mkstemp
 from time import sleep
-from typing import Optional
+from typing import Any, Optional
 from unittest import TestCase
 
 from loctrkd.common import init_protocols
@@ -117,6 +121,28 @@ class TestWithServers(TestCase):
                 pass
 
 
+class Fuzz(TestWithServers):
+    def setUp(self, *args: str, **kwargs: Any) -> None:
+        super().setUp("collector")
+        self.rnd = Random()
+        for fam, typ, pro, cnm, skadr in getaddrinfo(
+            "127.0.0.1",
+            self.conf.getint("collector", "port"),
+            family=AF_INET,
+            type=SOCK_STREAM,
+        ):
+            break  # Just take the first element
+        self.sock = socket(AF_INET, SOCK_STREAM)
+        self.sock.connect(skadr)
+
+    def tearDown(self) -> None:
+        sleep(1)  # give collector some time
+        send_and_drain(self.sock, None)
+        self.sock.close()
+        sleep(1)  # Let the server close their side
+        super().tearDown()
+
+
 def send_and_drain(sock: SocketType, buf: Optional[bytes]) -> None:
     if buf is not None:
         sock.send(buf)
diff --git a/test/test_fuzz.py b/test/test_fuzz.py
deleted file mode 100644
index ecedb2a..0000000
--- a/test/test_fuzz.py
+++ /dev/null
@@ -1,48 +0,0 @@
-""" Send junk to the collector """
-
-from random import Random
-from socket import getaddrinfo, socket, AF_INET, SOCK_STREAM
-from time import sleep
-from typing import Any
-import unittest
-from .common import send_and_drain, TestWithServers
-
-REPEAT: int = 1000000
-
-
-class Fuzz(TestWithServers):
-    def setUp(self, *args: str, **kwargs: Any) -> None:
-        super().setUp("collector")
-        self.rnd = Random()
-        for fam, typ, pro, cnm, skadr in getaddrinfo(
-            "127.0.0.1",
-            self.conf.getint("collector", "port"),
-            family=AF_INET,
-            type=SOCK_STREAM,
-        ):
-            break  # Just take the first element
-        self.sock = socket(AF_INET, SOCK_STREAM)
-        self.sock.connect(skadr)
-
-    def tearDown(self) -> None:
-        sleep(1)  # give collector some time
-        send_and_drain(self.sock, None)
-        self.sock.close()
-        sleep(1)  # Let the server close their side
-        super().tearDown()
-
-    def test_stream(self) -> None:
-        for _ in range(REPEAT):
-            size = self.rnd.randint(1, 5000)
-            buf = self.rnd.randbytes(size)
-            send_and_drain(self.sock, buf)
-
-    def test_msgs(self) -> None:
-        for _ in range(REPEAT):
-            size = self.rnd.randint(0, 300)
-            buf = b"xx" + self.rnd.randbytes(size) + b"\r\n"
-            send_and_drain(self.sock, buf)
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/test/test_fuzz_stream.py b/test/test_fuzz_stream.py
new file mode 100644
index 0000000..64c5a9f
--- /dev/null
+++ b/test/test_fuzz_stream.py
@@ -0,0 +1,18 @@
+""" Send junk to the collector """
+
+import unittest
+from .common import send_and_drain, TestWithServers, Fuzz
+
+REPEAT: int = 1000000
+
+
+class FuzzStream(Fuzz):
+    def test_stream(self) -> None:
+        for _ in range(REPEAT):
+            size = self.rnd.randint(1, 5000)
+            buf = self.rnd.randbytes(size)
+            send_and_drain(self.sock, buf)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/test_fuzz_zx303.py b/test/test_fuzz_zx303.py
new file mode 100644
index 0000000..4878f0e
--- /dev/null
+++ b/test/test_fuzz_zx303.py
@@ -0,0 +1,18 @@
+""" Send junk to the collector """
+
+import unittest
+from .common import send_and_drain, TestWithServers, Fuzz
+
+REPEAT: int = 1000000
+
+
+class FuzzMsgs(Fuzz):
+    def test_msgs(self) -> None:
+        for _ in range(REPEAT):
+            size = self.rnd.randint(0, 300)
+            buf = b"xx" + self.rnd.randbytes(size) + b"\r\n"
+            send_and_drain(self.sock, buf)
+
+
+if __name__ == "__main__":
+    unittest.main()
-- 
2.39.2