Improve tests

pytroll · May 21, 2024 · 6364282 · 6364282
1 parent c7b23ed
commit 6364282
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 38 deletions.
diff --git a/posttroll/address_receiver.py b/posttroll/address_receiver.py
@@ -159,23 +159,13 @@ def _run(self):
             try:
                 while self._do_run:
                     try:
-                        rerun = True
-                        while rerun:
-                            try:
-                                data, fromaddr = recv()
-                                rerun = False
-                            except TimeoutError:
-                                if self._do_run:
-                                    continue
-                                else:
-                                    raise
-                        if self._multicast_enabled:
-                            ip_, port = fromaddr
-                            if self._restrict_to_localhost and ip_ not in self._local_ips:
-                                # discard external message
-                                LOGGER.debug("Discard external message")
-                                continue
-                        LOGGER.debug("data %s", data)
+                        data, fromaddr = recv()
+                    except TimeoutError:
+                        if self._do_run:
+                            continue
+                        else:
+                            raise
+
                     except SocketTimeout:
                         if self._multicast_enabled:
                             LOGGER.debug("Multicast socket timed out on recv!")
@@ -186,6 +176,13 @@ def _run(self):
                         self._check_age(pub, min_interval=self._max_age / 20)
                         if self._do_heartbeat:
                             pub.heartbeat(min_interval=29)
+                    if self._multicast_enabled:
+                        ip_, port = fromaddr
+                        if self._restrict_to_localhost and ip_ not in self._local_ips:
+                            # discard external message
+                            LOGGER.debug("Discard external message")
+                            continue
+                    LOGGER.debug("data %s", data)
                     msg = Message.decode(data)
                     name = msg.subject.split("/")[1]
                     if msg.type == "info" and msg.subject.lower().startswith(self._subject):

diff --git a/posttroll/tests/test_nameserver.py b/posttroll/tests/test_nameserver.py
@@ -179,7 +179,7 @@ def test_pub_sub_add_rm(multicast_enabled):
                 for msg in sub.recv(.1):
                     if msg is None:
                         break
-                time.sleep(.3)
+                time.sleep(0.3)
                 assert len(sub.addresses) == 0
                 with Publish("data_provider_2", 0, ["another_data"], nameservers=nameservers):
                     time.sleep(.1)
@@ -225,23 +225,28 @@ def test_noisypublisher_heartbeat():
     from posttroll.publisher import NoisyPublisher
     from posttroll.subscriber import Subscribe
 
-    ns_ = NameServer()
-    thr = Thread(target=ns_.run)
-    thr.start()
+    min_interval = 10
 
-    pub = NoisyPublisher("test")
-    pub.start()
-    time.sleep(0.2)
-
-    with Subscribe("test", topics="/heartbeat/test", nameserver="localhost") as sub:
-        time.sleep(0.2)
-        pub.heartbeat(min_interval=10)
-        msg = next(sub.recv(1))
-    assert msg.type == "beat"
-    assert msg.data == {"min_interval": 10}
-    pub.stop()
-    ns_.stop()
-    thr.join()
+    try:
+        with config.set(address_publish_port=free_port(), nameserver_port=free_port()):
+            ns_ = NameServer()
+            thr = Thread(target=ns_.run)
+            thr.start()
+
+            pub = NoisyPublisher("test")
+            pub.start()
+            time.sleep(0.2)
+
+            with Subscribe("test", topics="/heartbeat/test", nameserver="localhost") as sub:
+                time.sleep(0.2)
+                pub.heartbeat(min_interval=min_interval)
+                msg = next(sub.recv(1))
+            assert msg.type == "beat"
+            assert msg.data == {"min_interval": min_interval}
+    finally:
+        pub.stop()
+        ns_.stop()
+        thr.join()
 
 
 def test_switch_backend_for_nameserver():

diff --git a/posttroll/tests/test_secure_zmq_backend.py b/posttroll/tests/test_secure_zmq_backend.py
@@ -153,10 +153,10 @@ def test_switch_to_secure_backend_for_nameserver(tmp_path):
     server_public_key_file, server_secret_key_file = zmq.auth.create_certificates(tmp_path, "server")
     client_public_key_file, client_secret_key_file = zmq.auth.create_certificates(tmp_path, "client")
     with config.set(backend="secure_zmq",
-                client_secret_key_file=client_secret_key_file,
-                clients_public_keys_directory=os.path.dirname(client_public_key_file),
-                server_public_key_file=server_public_key_file,
-                server_secret_key_file=server_secret_key_file):
+                    client_secret_key_file=client_secret_key_file,
+                    clients_public_keys_directory=os.path.dirname(client_public_key_file),
+                    server_public_key_file=server_public_key_file,
+                    server_secret_key_file=server_secret_key_file):
 
         with create_nameserver_instance():
             res = get_pub_address("some_name")