import os import shutil import hashlib import multiprocessing from logging import FileHandler from unittest import TestCase import config from hashtree import HashTree from client import Client, Connection as ClientConnection, DeniedConnection from server import Miniserver from . import RedirectedOutput config.logger.removeHandler(config.handler) handler = FileHandler("/tmp/morevna.log") handler.setFormatter(config.formatter) config.logger.addHandler(handler) config.batch_size.hash = 8 config.batch_size.data = 8 data_dir = os.path.join(config.directory, "src/tests/data") filename = os.path.join(data_dir, "test.img") def compare_files(f1, f2): with open(f1, mode="rb") as f: h2 = hashlib.sha256(f.read()).hexdigest() with open(f2, mode="rb") as f: h = hashlib.sha256(f.read()).hexdigest() return (h, h2) class TestMorevna(RedirectedOutput, TestCase): _stdout = None def setUp(self): src = os.path.join(data_dir, "test1.img") shutil.copyfile(src, filename) @classmethod def tearDownClass(cls): super().tearDownClass() os.remove(filename) def test_build(self): tree_file = os.path.join(data_dir, "test.bin") ref_file = os.path.join(data_dir, "test1.bin") tree = HashTree.from_file(os.path.join(data_dir, "test1.img")) tree.save(tree_file) self.assertEqual(*compare_files(ref_file, tree_file)) os.remove(tree_file) def test_push(self): config.port += 1 ms = Miniserver(filename) p = multiprocessing.Process(target=ms.serve) p.start() for client_file in ("test2.img", "test3.img", "test4.img"): client_file = os.path.join(data_dir, client_file) c = Client(client_file) with ClientConnection("127.0.0.1", config.port) as con: c.set_connection(con) c.init("push") blocks_to_transfer = c.negotiate() c.send_data(blocks_to_transfer) self.assertEqual(*compare_files(client_file, filename)) p.terminate() p.join() def test_pull(self): config.port += 1 server_file = os.path.join(data_dir, "test3.img") ms = Miniserver(server_file) p = multiprocessing.Process(target=ms.serve) p.start() c = Client(filename) with ClientConnection("127.0.0.1", config.port) as con: c.set_connection(con) c.init("pull") blocks_to_transfer = c.negotiate() c.pull_data(blocks_to_transfer) self.assertEqual(*compare_files(server_file, filename)) p.terminate() p.join() def test_deny(self): config.port += 1 ms = Miniserver(filename) p = multiprocessing.Process(target=ms.serve) p.start() c1 = Client(os.path.join(data_dir, "test2.img")) with ClientConnection("127.0.0.1", config.port) as con1: c1.set_connection(con1) c1.init("push") c2 = Client(os.path.join(data_dir, "test3.img")) with ClientConnection("127.0.0.1", config.port) as con2: c2.set_connection(con2) with self.assertRaises(DeniedConnection): c2.init("push") c1.send_data([]) # to unlock the server p.terminate() p.join()