From 21bcf742a4cd59dffdedbc367b708cef98d990c7 Mon Sep 17 00:00:00 2001
From: Stuckinaboot <aspyn@seas.upenn.edu>
Date: Tue, 2 Apr 2019 01:18:41 -0400
Subject: [PATCH] Brainstorming

---
 tests/pubsub/test_random_topology.py | 273 +++++++++++++++++++++++++++
 1 file changed, 273 insertions(+)
 create mode 100644 tests/pubsub/test_random_topology.py

diff --git a/tests/pubsub/test_random_topology.py b/tests/pubsub/test_random_topology.py
new file mode 100644
index 0000000..4b1e13a
--- /dev/null
+++ b/tests/pubsub/test_random_topology.py
@@ -0,0 +1,273 @@
+import asyncio
+import multiaddr
+import pytest
+import random
+
+from tests.utils import cleanup
+from libp2p import new_node
+from libp2p.peer.peerinfo import info_from_p2p_addr
+from libp2p.pubsub.pb import rpc_pb2
+from libp2p.pubsub.pubsub import Pubsub
+from libp2p.pubsub.floodsub import FloodSub
+from utils import generate_message_id, generate_RPC_packet
+
+# pylint: disable=too-many-locals
+
+async def connect(node1, node2):
+    """
+    Connect node1 to node2
+    """
+    addr = node2.get_addrs()[0]
+    info = info_from_p2p_addr(addr)
+    await node1.connect(info)
+
+async def perform_test_from_obj(obj):
+    """
+    Perform a floodsub test from a test obj.
+    test obj are composed as follows:
+    
+    {
+        "supported_protocols": ["supported/protocol/1.0.0",...],
+        "adj_list": {
+            "node1": ["neighbor1_of_node1", "neighbor2_of_node1", ...],
+            "node2": ["neighbor1_of_node2", "neighbor2_of_node2", ...],
+            ...
+        },
+        "topic_map": {
+            "topic1": ["node1_subscribed_to_topic1", "node2_subscribed_to_topic1", ...]
+        },
+        "messages": [
+            {
+                "topics": ["topic1_for_message", "topic2_for_message", ...],
+                "data": "some contents of the message (newlines are not supported)",
+                "node_id": "message sender node id"
+            },
+            ...
+        ]
+    }
+    NOTE: In adj_list, for any neighbors A and B, only list B as a neighbor of A
+    or B as a neighbor of A once. Do NOT list both A: ["B"] and B:["A"] as the behavior
+    is undefined (even if it may work)
+    """
+
+    # Step 1) Create graph
+    adj_list = obj["adj_list"]
+    node_map = {}
+    floodsub_map = {}
+    pubsub_map = {}
+
+    supported_protocols = obj["supported_protocols"]
+
+    tasks_connect = []
+    for start_node_id in adj_list:
+        # Create node if node does not yet exist
+        if start_node_id not in node_map:
+            node = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
+            await node.get_network().listen(multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0"))
+
+            node_map[start_node_id] = node
+
+            floodsub = FloodSub(supported_protocols)
+            floodsub_map[start_node_id] = floodsub
+            pubsub = Pubsub(node, floodsub, start_node_id)
+            pubsub_map[start_node_id] = pubsub
+
+        # For each neighbor of start_node, create if does not yet exist,
+        # then connect start_node to neighbor
+        for neighbor_id in adj_list[start_node_id]:
+            # Create neighbor if neighbor does not yet exist
+            if neighbor_id not in node_map:
+                neighbor_node = await new_node(transport_opt=["/ip4/127.0.0.1/tcp/0"])
+                await neighbor_node.get_network().listen(multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0"))
+                
+                node_map[neighbor_id] = neighbor_node
+
+                floodsub = FloodSub(supported_protocols)
+                floodsub_map[neighbor_id] = floodsub
+                pubsub = Pubsub(neighbor_node, floodsub, neighbor_id)
+                pubsub_map[neighbor_id] = pubsub
+
+            # Connect node and neighbor
+            # await connect(node_map[start_node_id], node_map[neighbor_id])
+            tasks_connect.append(asyncio.ensure_future(connect(node_map[start_node_id], node_map[neighbor_id])))
+    tasks_connect.append(asyncio.sleep(2))
+    await asyncio.gather(*tasks_connect)
+
+    # Allow time for graph creation before continuing
+    # await asyncio.sleep(0.25)
+
+    # Step 2) Subscribe to topics
+    queues_map = {}
+    topic_map = obj["topic_map"]
+
+    tasks_topic = []
+    tasks_topic_data = []
+    for topic in topic_map:
+        for node_id in topic_map[topic]:
+            """
+            # Subscribe node to topic
+            q = await pubsub_map[node_id].subscribe(topic)
+
+            # Create topic-queue map for node_id if one does not yet exist
+            if node_id not in queues_map:
+                queues_map[node_id] = {}
+
+            # Store queue in topic-queue map for node
+            queues_map[node_id][topic] = q
+            """
+            tasks_topic.append(asyncio.ensure_future(pubsub_map[node_id].subscribe(topic)))
+            tasks_topic_data.append((node_id, topic))
+    tasks_topic.append(asyncio.sleep(2))
+
+    # Gather is like Promise.all
+    responses = await asyncio.gather(*tasks_topic, return_exceptions=True)
+    for i in range(len(responses) - 1):
+        q = responses[i]
+        node_id, topic = tasks_topic_data[i]
+        if node_id not in queues_map:
+            queues_map[node_id] = {}
+
+        # Store queue in topic-queue map for node
+        queues_map[node_id][topic] = q
+
+    # Allow time for subscribing before continuing
+    # await asyncio.sleep(0.01)
+
+    # Step 3) Publish messages
+    topics_in_msgs_ordered = []
+    messages = obj["messages"]
+    tasks_publish = []
+    for msg in messages:
+        topics = msg["topics"]
+
+        data = msg["data"]
+        node_id = msg["node_id"]
+
+        # Get actual id for sender node (not the id from the test obj)
+        actual_node_id = str(node_map[node_id].get_id())
+
+        # Create correctly formatted message
+        msg_talk = generate_RPC_packet(actual_node_id, topics, data, generate_message_id())
+
+        # Publish message
+        # await floodsub_map[node_id].publish(actual_node_id, msg_talk.to_str())
+        tasks_publish.append(asyncio.ensure_future(floodsub_map[node_id].publish(\
+            actual_node_id, msg_talk.SerializeToString())))
+
+        # For each topic in topics, add topic, msg_talk tuple to ordered test list
+        # TODO: Update message sender to be correct message sender before
+        # adding msg_talk to this list
+        for topic in topics:
+            topics_in_msgs_ordered.append((topic, msg_talk))
+
+    # Allow time for publishing before continuing
+    # await asyncio.sleep(0.4)
+    tasks_publish.append(asyncio.sleep(2))
+    await asyncio.gather(*tasks_publish)
+
+    # Step 4) Check that all messages were received correctly.
+    # TODO: Check message sender too
+    for i in range(len(topics_in_msgs_ordered)):
+        topic, actual_msg = topics_in_msgs_ordered[i]
+        for node_id in topic_map[topic]:
+            # Get message from subscription queue
+            msg_on_node_str = await queues_map[node_id][topic].get()
+            assert actual_msg.publish[0].SerializeToString() == msg_on_node_str.SerializeToString()
+
+    # Success, terminate pending tasks.
+    await cleanup()
+
+# @pytest.mark.asyncio
+# async def test_simple_two_nodes_test_obj():
+#     test_obj = {
+#         "supported_protocols": ["/floodsub/1.0.0"],
+#         "adj_list": {
+#             "A": ["B"],
+#             "B": ["A"]
+#         },
+#         "topic_map": {
+#             "topic1": ["B"]
+#         },
+#         "messages": [
+#             {
+#                 "topics": ["topic1"],
+#                 "data": "foo",
+#                 "node_id": "A"
+#             }
+#         ]
+#     }
+#     await perform_test_from_obj(test_obj)
+
+def generate_random_topology(num_nodes, topic_density, num_topics, max_nodes_per_topic, max_msgs_per_topic):
+    nodes = range(num_nodes)
+    
+    # Create a separate graph for each topic
+    topic_graphs = {}
+    for topic in range(num_topics):
+        # TODO: Pick random num of nodes to be in topic (at least 1)
+        num_nodes_in_topic = max_nodes_per_topic
+        nodes_in_topic = random.sample(nodes, num_nodes_in_topic)
+        print("***Nodes in topic***")
+        print(num_nodes_in_topic)
+        print(nodes_in_topic)
+
+        # Create initial graph by connecting each node to its previous node
+        # This ensures the graph is connected
+        graph = {}
+
+        graph[nodes_in_topic[0]] = []
+
+        max_num_edges = num_nodes_in_topic * (num_nodes_in_topic - 1) / 2
+        num_edges = 1
+
+        for i in range(1, len(nodes_in_topic)):
+            prev = nodes_in_topic[i - 1]
+            curr = nodes_in_topic[i]
+
+            graph[curr] = [prev]
+            graph[prev].append(curr)
+            num_edges += 1
+
+        # Add random edges until topic density is hit
+        while num_edges / max_num_edges < topic_density:
+            selected_nodes = random.sample(nodes_in_topic, 2)
+
+            # Only add the nodes as neighbors if they are not already neighbors
+            if selected_nodes[0] not in graph[selected_nodes[1]]:
+                graph[selected_nodes[0]].append(selected_nodes[1])
+                graph[selected_nodes[1]].append(selected_nodes[0])
+                num_edges += 1
+
+        topic_graphs[topic] = graph
+        print(graph)
+
+    # Generate network graph from union of topic graphs
+    network_graph = {}
+
+    for topic in topic_graphs:
+        graph = topic_graphs[topic]
+        for node in graph:
+            # Add node if not in network graph
+            if node not in network_graph:
+                network_graph[node] = []
+            for neighbor in graph[node]:
+                # Add neighbor if not in network graph
+                if neighbor not in network_graph:
+                    network_graph[neighbor] = []
+
+                # Add edge if not in network graph
+                if neighbor not in network_graph[node]:
+                    network_graph[node].append(neighbor)
+                    network_graph[neighbor].append(node)
+
+def test_simple_random():
+    num_nodes = 4
+    topic_density = 0.5
+    num_topics = 2
+    max_nodes_per_topic = 4
+    max_msgs_per_topic = 1
+    topology = generate_random_topology(num_nodes, topic_density, num_topics,\
+        max_nodes_per_topic, max_msgs_per_topic)
+    print("TOPOLOGY")
+    print(topology)
+