py-libp2p/tests/pubsub/dummy_account_node.py
mhchia c72dfe1dd3
Use factories and fixtures in pubsub tests
Done
- Add factories using factory-boy
- Modify fixtures and tests to use factories
- Modify tests to use fixtures and factories
- Clean up
2019-08-01 13:30:08 +08:00

133 lines
4.4 KiB
Python

import asyncio
import uuid
from libp2p.host.host_interface import IHost
from libp2p.pubsub.floodsub import FloodSub
from libp2p.pubsub.pubsub import Pubsub
from tests.configs import LISTEN_MADDR
from .factories import FloodsubFactory, PubsubFactory
from .utils import message_id_generator
CRYPTO_TOPIC = "ethereum"
# Message format:
# Sending crypto: <source>,<dest>,<amount as integer>
# Ex. send,aspyn,alex,5
# Set crypto: <dest>,<amount as integer>
# Ex. set,rob,5
# Determine message type by looking at first item before first comma
class DummyAccountNode:
"""
Node which has an internal balance mapping, meant to serve as
a dummy crypto blockchain. There is no actual blockchain, just a simple
map indicating how much crypto each user in the mappings holds
"""
libp2p_node: IHost
pubsub: Pubsub
floodsub: FloodSub
def __init__(self, libp2p_node: IHost, pubsub: Pubsub, floodsub: FloodSub):
self.libp2p_node = libp2p_node
self.pubsub = pubsub
self.floodsub = floodsub
self.balances = {}
self.next_msg_id_func = message_id_generator(0)
self.node_id = str(uuid.uuid1())
@classmethod
async def create(cls):
"""
Create a new DummyAccountNode and attach a libp2p node, a floodsub, and a pubsub
instance to this new node
We use create as this serves as a factory function and allows us
to use async await, unlike the init function
"""
pubsub = PubsubFactory(router=FloodsubFactory())
await pubsub.host.get_network().listen(LISTEN_MADDR)
return cls(libp2p_node=pubsub.host, pubsub=pubsub, floodsub=pubsub.router)
async def handle_incoming_msgs(self):
"""
Handle all incoming messages on the CRYPTO_TOPIC from peers
"""
while True:
incoming = await self.q.get()
msg_comps = incoming.data.decode("utf-8").split(",")
if msg_comps[0] == "send":
self.handle_send_crypto(msg_comps[1], msg_comps[2], int(msg_comps[3]))
elif msg_comps[0] == "set":
self.handle_set_crypto(msg_comps[1], int(msg_comps[2]))
async def setup_crypto_networking(self):
"""
Subscribe to CRYPTO_TOPIC and perform call to function that handles
all incoming messages on said topic
"""
self.q = await self.pubsub.subscribe(CRYPTO_TOPIC)
asyncio.ensure_future(self.handle_incoming_msgs())
async def publish_send_crypto(self, source_user, dest_user, amount):
"""
Create a send crypto message and publish that message to all other nodes
:param source_user: user to send crypto from
:param dest_user: user to send crypto to
:param amount: amount of crypto to send
"""
msg_contents = "send," + source_user + "," + dest_user + "," + str(amount)
await self.pubsub.publish(CRYPTO_TOPIC, msg_contents.encode())
async def publish_set_crypto(self, user, amount):
"""
Create a set crypto message and publish that message to all other nodes
:param user: user to set crypto for
:param amount: amount of crypto
"""
msg_contents = "set," + user + "," + str(amount)
await self.pubsub.publish(CRYPTO_TOPIC, msg_contents.encode())
def handle_send_crypto(self, source_user, dest_user, amount):
"""
Handle incoming send_crypto message
:param source_user: user to send crypto from
:param dest_user: user to send crypto to
:param amount: amount of crypto to send
"""
if source_user in self.balances:
self.balances[source_user] -= amount
else:
self.balances[source_user] = -amount
if dest_user in self.balances:
self.balances[dest_user] += amount
else:
self.balances[dest_user] = amount
def handle_set_crypto(self, dest_user, amount):
"""
Handle incoming set_crypto message
:param dest_user: user to set crypto for
:param amount: amount of crypto
"""
self.balances[dest_user] = amount
def get_balance(self, user):
"""
Get balance in crypto for a particular user
:param user: user to get balance for
:return: balance of user
"""
if user in self.balances:
return self.balances[user]
else:
return -1