diff --git a/src/onionrblocks/storagecounter.py b/src/onionrblocks/storagecounter.py index bf919d69..ef35f64f 100755 --- a/src/onionrblocks/storagecounter.py +++ b/src/onionrblocks/storagecounter.py @@ -1,8 +1,17 @@ """ - Onionr - Private P2P Communication +Onionr - Private P2P Communication. - Keeps track of how much disk space we're using +Keep track of how much disk space we're using """ +from pathlib import Path + +from threading import Thread + +from watchdog.observers import Observer +from watchdog.events import FileSystemEventHandler + +import config +from filepaths import usage_file """ This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by @@ -17,17 +26,46 @@ You should have received a copy of the GNU General Public License along with this program. If not, see . """ -import config, filepaths config.reload() + + +def _read_data_file(f) -> int: + amount = 0 + try: + with open(f, 'r') as f: + amount = int(f.read()) + except FileNotFoundError: + pass + except ValueError: + pass # Possibly happens when the file is empty + return amount + + class StorageCounter: def __init__(self): - self.data_file = filepaths.usage_file - return + self.data_file = usage_file + self.amount: int = None + Path(self.data_file).touch() - def is_full(self)->bool: - """Returns if the allocated disk space is full (this is Onionr config, not true FS capacity)""" + def auto_refresher(): + class Refresher(FileSystemEventHandler): + @staticmethod + def on_modified(event): + self.amount = _read_data_file(self.data_file) + observer = Observer() + observer.schedule(Refresher(), usage_file) + observer.start() + while observer.is_alive(): + # call import func with timeout + observer.join(120) + Thread(target=auto_refresher, daemon=True).start() + self.amount = _read_data_file(self.data_file) + + def is_full(self) -> bool: + """Returns if the allocated disk space is full (this is Onionr config, + not true FS capacity)""" ret_data = False - if config.get('allocations.disk', 1073741824) <= (self.get_amount() + 1000): + if config.get('allocations.disk', 1073741824) <= (self.amount + 1000): ret_data = True return ret_data @@ -35,26 +73,15 @@ class StorageCounter: with open(self.data_file, 'w') as data_file: data_file.write(str(data)) - def get_amount(self)->int: - """Return how much disk space we're using (according to record)""" - ret_data = 0 - try: - with open(self.data_file, 'r') as data_file: - ret_data = int(data_file.read()) - except FileNotFoundError: - pass - except ValueError: - pass # Possibly happens when the file is empty - return ret_data - - def get_percent(self)->int: + def get_percent(self) -> int: """Return percent (decimal/float) of disk space we're using""" - amount = self.get_amount() + amount = self.amount return round(amount / config.get('allocations.disk', 2000000000), 2) - def add_bytes(self, amount)->int: - """Record that we are now using more disk space, unless doing so would exceed configured max""" - new_amount = amount + self.get_amount() + def add_bytes(self, amount) -> int: + """Record that we are now using more disk space, + unless doing so would exceed configured max""" + new_amount = amount + self.amount ret_data = new_amount if new_amount > config.get('allocations.disk', 2000000000): ret_data = False @@ -62,8 +89,8 @@ class StorageCounter: self._update(new_amount) return ret_data - def remove_bytes(self, amount)->int: + def remove_bytes(self, amount) -> int: """Record that we are now using less disk space""" - new_amount = self.get_amount() - amount + new_amount = self.amount - amount self._update(new_amount) return new_amount diff --git a/tests/test_storagecounter.py b/tests/test_storagecounter.py index 8e6fdcca..1b2d2369 100644 --- a/tests/test_storagecounter.py +++ b/tests/test_storagecounter.py @@ -2,6 +2,7 @@ import sys, os sys.path.append(".") sys.path.append("src/") import unittest, uuid +from time import sleep TEST_DIR = 'testdata/%s-%s' % (uuid.uuid4(), os.path.basename(__file__)) + '/' os.environ["ONIONR_HOME"] = TEST_DIR @@ -29,7 +30,7 @@ class TestStorageCounter(unittest.TestCase): self.assertIsNotNone(config.get('allocations.disk')) self.assertGreaterEqual(config.get('allocations.disk'), 1000000) - + def test_insert_too_much(self): _test_setup() config.set('allocations.disk', 1000) @@ -38,11 +39,13 @@ class TestStorageCounter(unittest.TestCase): def test_count(self): _test_setup() counter = storagecounter.StorageCounter() - start_value = counter.get_amount() + start_value = counter.amount b_hash = onionrblocks.insert("test") - self.assertGreater(counter.get_amount(), start_value) + sleep(0.1) + self.assertGreater(counter.amount, start_value) onionrstorage.removeblock.remove_block(b_hash) - self.assertEqual(counter.get_amount(), start_value) - + sleep(0.1) + self.assertEqual(counter.amount, start_value) + unittest.main()