From 8ab387ac05ff36056c8b30571c9fbab0e0333713 Mon Sep 17 00:00:00 2001
From: zixuanzh <zixuanzh@seas.upenn.edu>
Date: Sun, 28 Apr 2019 15:00:43 -0400
Subject: [PATCH] refactor add_router in swarm

---
 libp2p/__init__.py      | 9 +++++----
 libp2p/network/swarm.py | 6 +++---
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/libp2p/__init__.py b/libp2p/__init__.py
index 3e133b64..9d779f9d 100644
--- a/libp2p/__init__.py
+++ b/libp2p/__init__.py
@@ -52,8 +52,8 @@ def initialize_default_kademlia_router(
 
 
 def initialize_default_swarm(
-        id_opt=None, transport_opt=None,
-        muxer_opt=None, sec_opt=None, peerstore_opt=None):
+        id_opt=None, transport_opt=None, muxer_opt=None,
+        sec_opt=None, peerstore_opt=None, disc_opt=None):
     """
     initialize swarm when no swarm is passed in
     :param id_opt: optional id for host
@@ -79,7 +79,8 @@ def initialize_default_swarm(
     upgrader = TransportUpgrader(sec, muxer)
 
     peerstore = peerstore_opt or PeerStore()
-    swarm_opt = Swarm(id_opt, peerstore, upgrader, transport)
+    swarm_opt = Swarm(id_opt, peerstore,\
+                      upgrader, transport, disc_opt)
 
     return swarm_opt
 
@@ -106,7 +107,7 @@ async def new_node(
         swarm_opt = initialize_default_swarm(
             id_opt=id_opt, transport_opt=transport_opt,
             muxer_opt=muxer_opt, sec_opt=sec_opt,
-            peerstore_opt=peerstore_opt)
+            peerstore_opt=peerstore_opt, disc_opt=disc_opt)
 
     swarm_opt.add_router(disc_opt)
 
diff --git a/libp2p/network/swarm.py b/libp2p/network/swarm.py
index 8e6ba732..9b37b50a 100644
--- a/libp2p/network/swarm.py
+++ b/libp2p/network/swarm.py
@@ -10,17 +10,17 @@ from .stream.net_stream import NetStream
 from .connection.raw_connection import RawConnection
 
 class Swarm(INetwork):
-    # pylint: disable=too-many-instance-attributes, cell-var-from-loop
+    # pylint: disable=too-many-instance-attributes,cell-var-from-loop,too-many-arguments
 
-    def __init__(self, peer_id, peerstore, upgrader, transport):
+    def __init__(self, peer_id, peerstore, upgrader, transport, router):
         self.self_id = peer_id
         self.peerstore = peerstore
         self.upgrader = upgrader
         self.transport = transport
+        self.router = router
         self.connections = dict()
         self.listeners = dict()
         self.stream_handlers = dict()
-        self.router = None
 
         # Protocol muxing
         self.multiselect = Multiselect()