Add artificial delay for peer.

p2pfs/core/peer.py

@@ -23,6 +23,8 @@ class Peer(MessageServer):
 
         self._pending_publish = set()
 
+        self._delay = 0
+
     async def start(self):
         # connect to server
         try:
@@ -53,6 +55,9 @@ class Peer(MessageServer):
             pass
         await super().stop()
 
+    def set_delay(self, delay):
+        self._delay = 0 if delay is None else delay
+
     async def publish(self, local_file, remote_name=None):
         if not os.path.exists(local_file):
             return False, 'File {} doesn\'t exist'.format(local_file)
@@ -190,6 +195,9 @@ class Peer(MessageServer):
             message = await self._read_message(reader)
             if message is None:
                 break
+            # artificial delay for peer
+            if self._delay != 0:
+                await asyncio.sleep(self._delay)
             message_type = MessageType(message['type'])
             if message_type == MessageType.PEER_REQUEST_CHUNK:
                 assert message['filename'] in self._file_map, 'File {} requested does not exist'.format(message['filename'])

p2pfs/ui/terminal.py

@@ -55,6 +55,13 @@ class PeerTerminal(aiocmd.Cmd):
         _, message = await self._peer.publish(arg)
         print(message)
 
+    async def do_set_delay(self, arg):
+        arg = arg.split(' ')[0]
+        if arg == '':
+            print('delay is required.')
+        else:
+            self._peer.set_delay(float(arg))
+
     async def do_list_files(self, arg):
         file_list_dict = await self._peer.list_file()
         table = BeautifulTable()

tests/test_core.py

@@ -2,6 +2,7 @@ import os
 import asyncio
 import pytest
 from p2pfs import Peer, Tracker
+import time
 from tests.conftest import fmd5, TEST_SMALL_FILE, TEST_LARGE_FILE, TEST_SMALL_FILE_SIZE, TEST_LARGE_FILE_SIZE
 
 pytestmark = pytest.mark.asyncio
@@ -123,6 +124,42 @@ async def test_download(unused_tcp_port):
     await asyncio.gather(*[peer.stop() for peer in peers])
 
 
+async def test_delay(unused_tcp_port):
+    tracker = Tracker('localhost', unused_tcp_port)
+    peers = tuple(Peer('localhost', 0, 'localhost', unused_tcp_port) for _ in range(2))
+    tracker_started = await tracker.start()
+    peer_started = await asyncio.gather(*[peer.start() for peer in peers])
+    assert tracker_started and peer_started
+
+    # peer0 publishes a small_file and peer1 publishes a large file
+    is_success, _ = await peers[0].publish(TEST_SMALL_FILE)
+    assert is_success
+    file_list = tracker.file_list()
+    assert TEST_SMALL_FILE in file_list
+    assert file_list[TEST_SMALL_FILE]['size'] == TEST_SMALL_FILE_SIZE
+    file_list = await peers[1].list_file()
+    assert TEST_SMALL_FILE in file_list
+
+    # download small file
+    start = time.time()
+    result, msg = await peers[1].download(TEST_SMALL_FILE, 'downloaded_' + TEST_SMALL_FILE)
+    assert result is True
+    assert os.path.exists('downloaded_' + TEST_SMALL_FILE)
+    assert fmd5(TEST_SMALL_FILE) == fmd5('downloaded_' + TEST_SMALL_FILE)
+    os.remove('downloaded_' + TEST_SMALL_FILE)
+    download_time = time.time() - start
+    start = time.time()
+    peers[0].set_delay(1)
+    result, msg = await peers[1].download(TEST_SMALL_FILE, 'downloaded_' + TEST_SMALL_FILE)
+    assert result is True
+    download_time_with_delay = time.time() - start
+    assert download_time_with_delay > download_time
+    peers[0].set_delay(0)
+
+    await tracker.stop()
+    await asyncio.gather(*[peer.stop() for peer in peers])
+
+
 async def test_peer_disconnect(unused_tcp_port):
     tracker = Tracker('localhost', unused_tcp_port)
     peer = Peer('localhost', 0, 'localhost', unused_tcp_port)

tests/test_terminal.py

@@ -29,6 +29,7 @@ async def test_terminals(unused_tcp_port, capsys):
     await tracker_terminal.do_list_files('')
     out, _ = capsys.readouterr()
     assert TEST_SMALL_FILE in out
+    await peer_terminals[1].do_set_delay('0')
     await peer_terminals[1].do_download(TEST_SMALL_FILE + ' ' + 'downloaded_' + TEST_SMALL_FILE)
     assert os.path.exists('downloaded_' + TEST_SMALL_FILE)
     assert fmd5(TEST_SMALL_FILE) == fmd5('downloaded_' + TEST_SMALL_FILE)