111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
import os
|
|
import hashlib
|
|
import asyncio
|
|
import uvloop
|
|
from p2pfs import Peer, Tracker
|
|
from tests.conftest import TEST_SMALL_FILE, TEST_LARGE_FILE, TEST_SMALL_FILE_SIZE, TEST_LARGE_FILE_SIZE
|
|
|
|
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
|
|
|
loop = asyncio.get_event_loop()
|
|
peers = tuple(Peer('localhost', 0, 'localhost', 8880) for _ in range(3))
|
|
tracker = Tracker('localhost', 8880)
|
|
|
|
|
|
def teardown_module(module):
|
|
""" stop the peers and tracker"""
|
|
loop.run_until_complete(asyncio.wait({peer.stop() for peer in peers}))
|
|
loop.run_until_complete(tracker.stop())
|
|
loop.close()
|
|
|
|
|
|
def fmd5(fname):
|
|
""" calculate the md5 value of a file
|
|
:param fname: file name
|
|
:return: md5 value.
|
|
"""
|
|
hash_md5 = hashlib.md5()
|
|
with open(fname, "rb") as f:
|
|
for chunk in iter(lambda: f.read(4096), b""):
|
|
hash_md5.update(chunk)
|
|
return hash_md5.hexdigest()
|
|
|
|
|
|
def test_server_refused():
|
|
peer = Peer('localhost', 0, 'localhost', 8880)
|
|
started = asyncio.get_event_loop().run_until_complete(peer.start())
|
|
assert not started
|
|
|
|
|
|
def test_start():
|
|
tracker_started = loop.run_until_complete(tracker.start())
|
|
# spawn peers concurrently
|
|
peers_started = \
|
|
loop.run_until_complete(asyncio.gather(*[peer.start() for peer in peers]))
|
|
assert tracker_started and all(peers_started)
|
|
|
|
|
|
def test_publish_refuse():
|
|
with open('test_publish_refuse', 'wb') as fout:
|
|
fout.write(os.urandom(100))
|
|
|
|
is_success, _ = loop.run_until_complete(peers[0].publish('test_publish_refuse'))
|
|
assert is_success
|
|
is_success, _ = loop.run_until_complete(peers[0].publish('test_publish_refuse'))
|
|
assert not is_success
|
|
|
|
|
|
def test_publish_small():
|
|
# peer1 publish small file and peer2 downloads it
|
|
is_success, _ = loop.run_until_complete(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 = loop.run_until_complete(peers[1].list_file())
|
|
assert TEST_SMALL_FILE in file_list
|
|
|
|
|
|
def test_download_small():
|
|
def reporthook(chunk_num, chunk_size, total_size):
|
|
reporthook.value = (chunk_num, total_size)
|
|
|
|
# download small file
|
|
result, msg = loop.run_until_complete(
|
|
peers[1].download(TEST_SMALL_FILE, 'downloaded_' + TEST_SMALL_FILE, reporthook=reporthook)
|
|
)
|
|
assert result is True
|
|
assert os.path.exists('downloaded_' + TEST_SMALL_FILE)
|
|
assert fmd5(TEST_SMALL_FILE) == fmd5('downloaded_' + TEST_SMALL_FILE)
|
|
assert reporthook.value == (1, 1000)
|
|
os.remove('downloaded_' + TEST_SMALL_FILE)
|
|
|
|
|
|
def test_publish_large():
|
|
# publish big file
|
|
is_success, _ = loop.run_until_complete(peers[1].publish(TEST_LARGE_FILE))
|
|
assert is_success
|
|
file_list = tracker.file_list()
|
|
assert TEST_LARGE_FILE in file_list
|
|
assert file_list[TEST_LARGE_FILE]['size'] == TEST_LARGE_FILE_SIZE
|
|
file_list = loop.run_until_complete(peers[0].list_file())
|
|
assert TEST_LARGE_FILE in file_list
|
|
|
|
|
|
def test_download_large_single():
|
|
# download from single peer
|
|
result, msg = loop.run_until_complete(peers[0].download(TEST_LARGE_FILE, 'downloaded_' + TEST_LARGE_FILE))
|
|
assert result is True
|
|
assert os.path.exists('downloaded_' + TEST_LARGE_FILE)
|
|
assert fmd5(TEST_LARGE_FILE) == fmd5('downloaded_' + TEST_LARGE_FILE)
|
|
os.remove('downloaded_' + TEST_LARGE_FILE)
|
|
|
|
|
|
def test_download_large_multi():
|
|
# test download from multiple peers
|
|
result, msg = loop.run_until_complete(peers[2].download(TEST_LARGE_FILE,
|
|
'downloaded_'+ TEST_LARGE_FILE + '_from_multiple_peers'))
|
|
assert result is True
|
|
assert os.path.exists('downloaded_' + TEST_LARGE_FILE + '_from_multiple_peers')
|
|
assert fmd5(TEST_LARGE_FILE) == fmd5('downloaded_' + TEST_LARGE_FILE + '_from_multiple_peers')
|
|
os.remove('downloaded_' + TEST_LARGE_FILE + '_from_multiple_peers')
|