From 828ae69c66c8c6fe6964c98f1bfb5624a4ff0501 Mon Sep 17 00:00:00 2001
From: NIC619 <twedusuck@gmail.com>
Date: Tue, 26 Nov 2019 11:52:31 +0800
Subject: [PATCH] Apply PR feedback: `add_key_pair`

---
 libp2p/__init__.py                 |  3 +--
 libp2p/peer/peerstore.py           | 13 +++++++++++--
 libp2p/peer/peerstore_interface.py | 10 +++++++++-
 libp2p/tools/factories.py          |  3 +--
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/libp2p/__init__.py b/libp2p/__init__.py
index 75b67644..291b3a52 100644
--- a/libp2p/__init__.py
+++ b/libp2p/__init__.py
@@ -106,8 +106,7 @@ def initialize_default_swarm(
 
     peerstore = peerstore_opt or PeerStore()
     # Store our key pair in peerstore
-    peerstore.add_pubkey(id_opt, key_pair.public_key)
-    peerstore.add_privkey(id_opt, key_pair.private_key)
+    peerstore.add_key_pair(id_opt, key_pair)
 
     # TODO: Initialize discovery if not presented
     return Swarm(id_opt, peerstore, upgrader, transport)
diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py
index 4cff54e2..c011e3ef 100644
--- a/libp2p/peer/peerstore.py
+++ b/libp2p/peer/peerstore.py
@@ -2,7 +2,7 @@ from typing import Any, Dict, List, Optional, Sequence, Set
 
 from multiaddr import Multiaddr
 
-from libp2p.crypto.keys import PrivateKey, PublicKey
+from libp2p.crypto.keys import KeyPair, PrivateKey, PublicKey
 
 from .id import ID
 from .peerdata import PeerData, PeerDataError
@@ -34,7 +34,7 @@ class PeerStore(IPeerStore):
             return self.peer_data_map[peer_id]
         data = PeerData()
         self.peer_data_map[peer_id] = data
-        return self.peer_data_map[peer_id]
+        return data
 
     def peer_info(self, peer_id: ID) -> Optional[PeerInfo]:
         """
@@ -196,6 +196,15 @@ class PeerStore(IPeerStore):
             return self.peer_privkey_map[peer_id]
         raise PeerStoreError("peer ID not found")
 
+    def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None:
+        """
+        :param peer_id: peer ID to add private key for
+        :param key_pair:
+        :raise PeerStoreError: if peer ID already has pubkey or privkey set
+        """
+        self.add_pubkey(peer_id, key_pair.public_keypubkey)
+        self.add_privkey(peer_id, key_pair.private_key)
+
     def peers_with_keys(self) -> Set[ID]:
         """
         :return: all of the peer IDs which has pubkey/privkey stored in peer store
diff --git a/libp2p/peer/peerstore_interface.py b/libp2p/peer/peerstore_interface.py
index 8ebe1cc6..618bbf5e 100644
--- a/libp2p/peer/peerstore_interface.py
+++ b/libp2p/peer/peerstore_interface.py
@@ -3,7 +3,7 @@ from typing import Any, List, Sequence, Set
 
 from multiaddr import Multiaddr
 
-from libp2p.crypto.keys import PrivateKey, PublicKey
+from libp2p.crypto.keys import KeyPair, PrivateKey, PublicKey
 
 from .addrbook_interface import IAddrBook
 from .id import ID
@@ -131,6 +131,14 @@ class IPeerStore(IAddrBook, IPeerMetadata):
         :raise PeerStoreError: if peer ID not found
         """
 
+    @abstractmethod
+    def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None:
+        """
+        :param peer_id: peer ID to add private key for
+        :param key_pair:
+        :raise PeerStoreError: if peer ID already has pubkey or privkey set
+        """
+
     @abstractmethod
     def peers_with_keys(self) -> Set[ID]:
         """
diff --git a/libp2p/tools/factories.py b/libp2p/tools/factories.py
index 405e6ed6..b5c16b85 100644
--- a/libp2p/tools/factories.py
+++ b/libp2p/tools/factories.py
@@ -36,8 +36,7 @@ from .utils import connect, connect_swarm
 
 def initialize_peerstore_with_our_keypair(self_id: ID, key_pair: KeyPair) -> PeerStore:
     peer_store = PeerStore()
-    peer_store.add_pubkey(self_id, key_pair.public_key)
-    peer_store.add_privkey(self_id, key_pair.private_key)
+    peer_store.add_key_pair(self_id, key_pair)
     return peer_store