From d5d6962dce9b8e72b0309771ab451c3ef8b79c7b Mon Sep 17 00:00:00 2001
From: NIC619 <twedusuck@gmail.com>
Date: Tue, 26 Nov 2019 16:12:50 +0800
Subject: [PATCH] Update Pubsub fixture and test

---
 libp2p/pubsub/pubsub.py     |  2 +-
 libp2p/tools/factories.py   |  1 +
 tests/pubsub/conftest.py    | 28 ++++++++++++++++++++++------
 tests/pubsub/test_pubsub.py | 15 +++++++++++++++
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/libp2p/pubsub/pubsub.py b/libp2p/pubsub/pubsub.py
index e49d0e8f..0fea734f 100644
--- a/libp2p/pubsub/pubsub.py
+++ b/libp2p/pubsub/pubsub.py
@@ -532,7 +532,7 @@ class Pubsub:
 
         # Check if signing is required and if so signature should be attached.
         if self.strict_signing:
-            if msg.signature == b'':
+            if msg.signature == b"":
                 logger.debug("Reject because no signature attached for msg: %s", msg)
                 return
             # Validate the signature of the message
diff --git a/libp2p/tools/factories.py b/libp2p/tools/factories.py
index b5c16b85..b189cfa1 100644
--- a/libp2p/tools/factories.py
+++ b/libp2p/tools/factories.py
@@ -153,6 +153,7 @@ class PubsubFactory(factory.Factory):
     router = None
     my_id = factory.LazyAttribute(lambda obj: obj.host.get_id())
     cache_size = None
+    strict_signing = False
 
 
 async def swarm_pair_factory(
diff --git a/tests/pubsub/conftest.py b/tests/pubsub/conftest.py
index 9dbe90b9..520fdf4b 100644
--- a/tests/pubsub/conftest.py
+++ b/tests/pubsub/conftest.py
@@ -4,14 +4,24 @@ from libp2p.tools.constants import GOSSIPSUB_PARAMS
 from libp2p.tools.factories import FloodsubFactory, GossipsubFactory, PubsubFactory
 
 
-def _make_pubsubs(hosts, pubsub_routers, cache_size):
+@pytest.fixture
+def is_strict_signing():
+    return False
+
+
+def _make_pubsubs(hosts, pubsub_routers, cache_size, is_strict_signing):
     if len(pubsub_routers) != len(hosts):
         raise ValueError(
             f"lenght of pubsub_routers={pubsub_routers} should be equaled to the "
             f"length of hosts={len(hosts)}"
         )
     return tuple(
-        PubsubFactory(host=host, router=router, cache_size=cache_size)
+        PubsubFactory(
+            host=host,
+            router=router,
+            cache_size=cache_size,
+            strict_signing=is_strict_signing,
+        )
         for host, router in zip(hosts, pubsub_routers)
     )
 
@@ -27,16 +37,22 @@ def gossipsub_params():
 
 
 @pytest.fixture
-def pubsubs_fsub(num_hosts, hosts, pubsub_cache_size):
+def pubsubs_fsub(num_hosts, hosts, pubsub_cache_size, is_strict_signing):
     floodsubs = FloodsubFactory.create_batch(num_hosts)
-    _pubsubs_fsub = _make_pubsubs(hosts, floodsubs, pubsub_cache_size)
+    _pubsubs_fsub = _make_pubsubs(
+        hosts, floodsubs, pubsub_cache_size, is_strict_signing
+    )
     yield _pubsubs_fsub
     # TODO: Clean up
 
 
 @pytest.fixture
-def pubsubs_gsub(num_hosts, hosts, pubsub_cache_size, gossipsub_params):
+def pubsubs_gsub(
+    num_hosts, hosts, pubsub_cache_size, gossipsub_params, is_strict_signing
+):
     gossipsubs = GossipsubFactory.create_batch(num_hosts, **gossipsub_params._asdict())
-    _pubsubs_gsub = _make_pubsubs(hosts, gossipsubs, pubsub_cache_size)
+    _pubsubs_gsub = _make_pubsubs(
+        hosts, gossipsubs, pubsub_cache_size, is_strict_signing
+    )
     yield _pubsubs_gsub
     # TODO: Clean up
diff --git a/tests/pubsub/test_pubsub.py b/tests/pubsub/test_pubsub.py
index ebe20037..48ef52bb 100644
--- a/tests/pubsub/test_pubsub.py
+++ b/tests/pubsub/test_pubsub.py
@@ -510,3 +510,18 @@ async def test_push_msg(pubsubs_fsub, monkeypatch):
     await pubsubs_fsub[0].push_msg(pubsubs_fsub[0].my_id, msg_2)
     await asyncio.sleep(0.01)
     assert not event.is_set()
+
+
+@pytest.mark.parametrize("num_hosts, is_strict_signing", ((2, True),))
+@pytest.mark.asyncio
+async def test_strict_signing(pubsubs_fsub, hosts, monkeypatch):
+    await connect(hosts[0], hosts[1])
+    await pubsubs_fsub[0].subscribe(TESTING_TOPIC)
+    await pubsubs_fsub[1].subscribe(TESTING_TOPIC)
+    await asyncio.sleep(1)
+
+    await pubsubs_fsub[0].publish(TESTING_TOPIC, TESTING_DATA)
+    await asyncio.sleep(1)
+
+    assert len(pubsubs_fsub[0].seen_messages) == 1
+    assert len(pubsubs_fsub[1].seen_messages) == 1