refactor(daa): remove global daa test mode

Author: glevco

hathor/builder/cli_builder.py

@@ -59,7 +59,7 @@ def check_or_raise(self, condition: bool, message: str) -> None:
     def create_manager(self, reactor: Reactor) -> HathorManager:
         import hathor
         from hathor.conf.get_settings import get_settings, get_settings_source
-        from hathor.daa import TestMode, _set_test_mode
+        from hathor.daa import TestMode
         from hathor.event.storage import EventMemoryStorage, EventRocksDBStorage, EventStorage
         from hathor.event.websocket.factory import EventWebsocketFactory
         from hathor.p2p.netfilter.utils import add_peer_id_blacklist
@@ -204,7 +204,13 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
             not_support_features=self._args.signal_not_support
         )
 
-        daa = DifficultyAdjustmentAlgorithm(settings=settings)
+        test_mode = TestMode.DISABLED
+        if self._args.test_mode_tx_weight:
+            test_mode = TestMode.TEST_TX_WEIGHT
+            if self.wallet:
+                self.wallet.test_mode = True
+
+        daa = DifficultyAdjustmentAlgorithm(settings=settings, test_mode=test_mode)
 
         vertex_verifiers = VertexVerifiers.create(settings=settings, daa=daa)
         verification_service = VerificationService(verifiers=vertex_verifiers)
@@ -273,11 +279,6 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
         if self._args.bootstrap:
             self.manager.add_peer_discovery(BootstrapPeerDiscovery(self._args.bootstrap))
 
-        if self._args.test_mode_tx_weight:
-            _set_test_mode(TestMode.TEST_TX_WEIGHT)
-            if self.wallet:
-                self.wallet.test_mode = True
-
         if self._args.x_rocksdb_indexes:
             self.log.warn('--x-rocksdb-indexes is now the default, no need to specify it')
             if self._args.memory_indexes:

hathor/daa.py

@@ -21,7 +21,7 @@
 
 from enum import IntFlag
 from math import log
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar, Optional
 
 from structlog import get_logger
 
@@ -45,26 +45,20 @@ class TestMode(IntFlag):
     TEST_ALL_WEIGHT = 3
 
 
-TEST_MODE = TestMode.DISABLED
-
-
-def _set_test_mode(mode: TestMode) -> None:
-    global TEST_MODE
-    logger.debug('change DAA test mode', from_mode=TEST_MODE.name, to_mode=mode.name)
-    TEST_MODE = mode
-
-
 class DifficultyAdjustmentAlgorithm:
+    singleton: ClassVar[Optional['DifficultyAdjustmentAlgorithm']] = None
 
-    def __init__(self, *, settings: HathorSettings) -> None:
+    def __init__(self, *, settings: HathorSettings, test_mode: TestMode = TestMode.DISABLED) -> None:
         self._settings = settings
         self.AVG_TIME_BETWEEN_BLOCKS = self._settings.AVG_TIME_BETWEEN_BLOCKS
         self.MIN_BLOCK_WEIGHT = self._settings.MIN_BLOCK_WEIGHT
+        self.TEST_MODE = test_mode
+        DifficultyAdjustmentAlgorithm.singleton = self
 
     @cpu.profiler(key=lambda _, block: 'calculate_block_difficulty!{}'.format(block.hash.hex()))
     def calculate_block_difficulty(self, block: 'Block') -> float:
         """ Calculate block weight according to the ascendents of `block`, using calculate_next_weight."""
-        if TEST_MODE & TestMode.TEST_BLOCK_WEIGHT:
+        if self.TEST_MODE & TestMode.TEST_BLOCK_WEIGHT:
             return 1.0
 
         if block.is_genesis:
@@ -79,7 +73,7 @@ def calculate_next_weight(self, parent_block: 'Block', timestamp: int) -> float:
 
         The weight must not be less than `MIN_BLOCK_WEIGHT`.
         """
-        if TEST_MODE & TestMode.TEST_BLOCK_WEIGHT:
+        if self.TEST_MODE & TestMode.TEST_BLOCK_WEIGHT:
             return 1.0
 
         from hathor.transaction import sum_weights
@@ -166,7 +160,7 @@ def minimum_tx_weight(self, tx: 'Transaction') -> float:
         """
         # In test mode we don't validate the minimum weight for tx
         # We do this to allow generating many txs for testing
-        if TEST_MODE & TestMode.TEST_TX_WEIGHT:
+        if self.TEST_MODE & TestMode.TEST_TX_WEIGHT:
             return 1.0
 
         if tx.is_genesis:

hathor/p2p/peer_id.py

@@ -27,9 +27,10 @@
 from twisted.internet.interfaces import ISSLTransport
 from twisted.internet.ssl import Certificate, CertificateOptions, TLSVersion, trustRootFromCertificates
 
-from hathor import daa
 from hathor.conf.get_settings import get_settings
+from hathor.daa import DifficultyAdjustmentAlgorithm
 from hathor.p2p.utils import connection_string_to_host, discover_dns, generate_certificate
+from hathor.util import not_none
 
 if TYPE_CHECKING:
     from hathor.p2p.protocol import HathorProtocol  # noqa: F401
@@ -344,7 +345,8 @@ def validate_entrypoint(self, protocol: 'HathorProtocol') -> Generator[Any, Any,
                     break
                 host = connection_string_to_host(entrypoint)
                 # TODO: don't use `daa.TEST_MODE` for this
-                result = yield discover_dns(host, daa.TEST_MODE)
+                test_mode = not_none(DifficultyAdjustmentAlgorithm.singleton).TEST_MODE
+                result = yield discover_dns(host, test_mode)
                 if protocol.connection_string in result:
                     # Found the entrypoint
                     found_entrypoint = True
@@ -363,7 +365,8 @@ def validate_entrypoint(self, protocol: 'HathorProtocol') -> Generator[Any, Any,
                 if connection_host == host:
                     found_entrypoint = True
                     break
-                result = yield discover_dns(host, daa.TEST_MODE)
+                test_mode = not_none(DifficultyAdjustmentAlgorithm.singleton).TEST_MODE
+                result = yield discover_dns(host, test_mode)
                 if connection_host in [connection_string_to_host(x) for x in result]:
                     # Found the entrypoint
                     found_entrypoint = True

hathor/simulator/simulator.py

@@ -22,7 +22,7 @@
 
 from hathor.builder import BuildArtifacts, Builder
 from hathor.conf.get_settings import get_settings
-from hathor.daa import DifficultyAdjustmentAlgorithm, TestMode, _set_test_mode
+from hathor.daa import DifficultyAdjustmentAlgorithm
 from hathor.manager import HathorManager
 from hathor.p2p.peer_id import PeerId
 from hathor.simulator.clock import HeapClock, MemoryReactorHeapClock
@@ -62,7 +62,6 @@ def _apply_patches(cls):
         Patches:
 
         - disable Transaction.resolve method
-        - set DAA test-mode to DISABLED (will actually run the pow function, that won't actually verify the pow)
         """
         from hathor.transaction import BaseTransaction
 
@@ -74,8 +73,6 @@ def resolve(self: BaseTransaction, update_time: bool = True) -> bool:
         cls._original_resolve = BaseTransaction.resolve
         BaseTransaction.resolve = resolve
 
-        _set_test_mode(TestMode.DISABLED)
-
     @classmethod
     def _remove_patches(cls):
         """ Remove the patches previously applied.

tests/p2p/test_split_brain.py

@@ -1,7 +1,7 @@
 import pytest
 from mnemonic import Mnemonic
 
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.graphviz import GraphvizVisualizer
 from hathor.simulator import FakeConnection
 from hathor.wallet import HDWallet
@@ -25,8 +25,8 @@ def create_peer(self, network, unlock_wallet=True):
         wallet = HDWallet(gap_limit=2)
         wallet._manually_initialize()
 
-        _set_test_mode(TestMode.TEST_ALL_WEIGHT)
         manager = super().create_peer(network, wallet=wallet)
+        manager.daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
         manager.avg_time_between_blocks = 64
 
         # Don't use it anywhere else. It is unsafe to generate mnemonic words like this.

tests/resources/base_resource.py

@@ -2,7 +2,7 @@
 from twisted.web import server
 from twisted.web.test.requesthelper import DummyRequest
 
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.util import json_dumpb, json_loadb
 from tests import unittest
 
@@ -19,7 +19,7 @@ def setUp(self, *, utxo_index: bool = False, unlock_wallet: bool = True) -> None
                 unlock_wallet=unlock_wallet
             )
             self.manager.allow_mining_without_peers()
-            _set_test_mode(TestMode.TEST_ALL_WEIGHT)
+            self.manager.daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
 
         def tearDown(self):
             return self.manager.stop()

tests/resources/transaction/test_create_tx.py

@@ -2,7 +2,7 @@
 
 from twisted.internet.defer import inlineCallbacks
 
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.transaction import Transaction
 from hathor.transaction.resources import CreateTxResource
 from hathor.transaction.scripts import P2PKH, create_base_script
@@ -195,7 +195,7 @@ def test_spend_tx_by_script(self):
 
     @inlineCallbacks
     def test_tx_propagate(self):
-        _set_test_mode(TestMode.DISABLED)      # disable test_mode so the weight is not 1
+        self.manager.daa.TEST_MODE = TestMode.DISABLED  # disable test_mode so the weight is not 1
         src_tx = self.unspent_tx
         output_address = 'HNXsVtRUmwDCtpcCJUrH4QiHo9kUKx199A'
         resp = (yield self.web.post('create_tx', {
@@ -233,7 +233,7 @@ def test_tx_propagate(self):
 
     @inlineCallbacks
     def test_tx_propagate_multiple_inputs(self):
-        _set_test_mode(TestMode.DISABLED)      # disable test_mode so the weight is not 1
+        self.manager.daa.TEST_MODE = TestMode.DISABLED  # disable test_mode so the weight is not 1
         output_address = 'HNXsVtRUmwDCtpcCJUrH4QiHo9kUKx199A'
         resp = (yield self.web.post('create_tx', {
             'inputs': [

tests/resources/wallet/test_send_tokens.py

@@ -2,7 +2,7 @@
 
 from twisted.internet.defer import inlineCallbacks
 
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.p2p.resources import MiningResource
 from hathor.wallet.resources import BalanceResource, HistoryResource, SendTokensResource
 from tests import unittest
@@ -168,7 +168,7 @@ def test_post(self):
 
     @inlineCallbacks
     def test_tx_weight(self):
-        _set_test_mode(TestMode.DISABLED)
+        self.manager.daa.TEST_MODE = TestMode.DISABLED
         add_new_blocks(self.manager, 3, advance_clock=1)
         add_blocks_unlock_reward(self.manager)
         self.reactor.advance(3)

tests/tx/test_blockchain.py

@@ -1,7 +1,7 @@
 from itertools import chain
 
 from hathor.conf import HathorSettings
-from hathor.daa import DifficultyAdjustmentAlgorithm, TestMode, _set_test_mode
+from hathor.daa import DifficultyAdjustmentAlgorithm, TestMode
 from hathor.transaction import sum_weights
 from hathor.transaction.storage import TransactionMemoryStorage
 from tests import unittest
@@ -390,8 +390,8 @@ def test_block_rewards(self):
 
     def test_daa_sanity(self):
         # sanity test the DAA
-        _set_test_mode(TestMode.DISABLED)
         manager = self.create_peer('testnet', tx_storage=self.tx_storage)
+        manager.daa.TEST_MODE = TestMode.DISABLED
         N = settings.BLOCK_DIFFICULTY_N_BLOCKS
         T = settings.AVG_TIME_BETWEEN_BLOCKS
         manager.avg_time_between_blocks = T
@@ -417,7 +417,7 @@ def test_daa_sanity(self):
             self.assertLess(new_weight, base_weight)
 
     def test_daa_weight_decay_amount(self):
-        _set_test_mode(TestMode.DISABLED)
+        self.daa.TEST_MODE = TestMode.DISABLED
         amount = settings.WEIGHT_DECAY_AMOUNT
 
         for distance in range(0, settings.WEIGHT_DECAY_ACTIVATE_DISTANCE, 10):
@@ -434,8 +434,8 @@ def test_daa_weight_decay_amount(self):
         self.assertAlmostEqual(self.daa.get_weight_decay_amount(distance), 11 * amount)
 
     def test_daa_weight_decay_blocks(self):
-        _set_test_mode(TestMode.DISABLED)
         manager = self.create_peer('testnet', tx_storage=self.tx_storage)
+        manager.daa.TEST_MODE = TestMode.DISABLED
         amount = settings.WEIGHT_DECAY_AMOUNT
 
         manager.daa.AVG_TIME_BETWEEN_BLOCKS = settings.AVG_TIME_BETWEEN_BLOCKS

tests/tx/test_cache_storage.py

@@ -1,4 +1,4 @@
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.transaction import Transaction, TransactionMetadata
 from hathor.transaction.storage import TransactionCacheStorage
 from tests import unittest
@@ -144,7 +144,7 @@ def test_flush_thread(self):
         self.cache_storage._flush_to_storage(self.cache_storage.dirty_txs.copy())
 
     def test_topological_sort_dfs(self):
-        _set_test_mode(TestMode.TEST_ALL_WEIGHT)
+        self.manager.daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
         add_new_blocks(self.manager, 11, advance_clock=1)
         tx = add_new_transactions(self.manager, 1, advance_clock=1)[0]
 

tests/tx/test_genesis.py

@@ -1,5 +1,5 @@
 from hathor.conf import HathorSettings
-from hathor.daa import DifficultyAdjustmentAlgorithm, TestMode, _set_test_mode
+from hathor.daa import DifficultyAdjustmentAlgorithm, TestMode
 from hathor.transaction.storage import TransactionMemoryStorage
 from hathor.verification.verification_service import VerificationService, VertexVerifiers
 from hathor.verification.vertex_verifier import VertexVerifier
@@ -69,10 +69,10 @@ def test_genesis_weight(self):
 
         # Validate the block and tx weight
         # in test mode weight is always 1
-        _set_test_mode(TestMode.TEST_ALL_WEIGHT)
+        self._daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
         self.assertEqual(self._daa.calculate_block_difficulty(genesis_block), 1)
         self.assertEqual(self._daa.minimum_tx_weight(genesis_tx), 1)
 
-        _set_test_mode(TestMode.DISABLED)
+        self._daa.TEST_MODE = TestMode.DISABLED
         self.assertEqual(self._daa.calculate_block_difficulty(genesis_block), genesis_block.weight)
         self.assertEqual(self._daa.minimum_tx_weight(genesis_tx), genesis_tx.weight)

tests/tx/test_tx.py

@@ -3,7 +3,7 @@
 from math import isinf, isnan
 
 from hathor.crypto.util import decode_address, get_address_from_public_key, get_private_key_from_bytes
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.transaction import MAX_OUTPUT_VALUE, Block, Transaction, TxInput, TxOutput
 from hathor.transaction.exceptions import (
     BlockWithInputs,
@@ -630,8 +630,8 @@ def test_update_timestamp(self):
         self.assertEquals(tx.timestamp, ts)
 
     def test_propagation_error(self):
-        _set_test_mode(TestMode.DISABLED)
         manager = self.create_peer('testnet', unlock_wallet=True)
+        manager.daa.TEST_MODE = TestMode.DISABLED
 
         # 1. propagate genesis
         genesis_block = self.genesis_blocks[0]

tests/tx/test_tx_storage.py

@@ -9,7 +9,7 @@
 from twisted.trial import unittest
 
 from hathor.conf import HathorSettings
-from hathor.daa import TestMode, _set_test_mode
+from hathor.daa import TestMode
 from hathor.transaction import Block, Transaction, TxInput, TxOutput
 from hathor.transaction.scripts import P2PKH
 from hathor.transaction.storage.exceptions import TransactionDoesNotExist
@@ -155,7 +155,7 @@ def test_storage_basic_v2(self):
         self.assertEqual(set(tx_parents_hash), {self.genesis_txs[0].hash, self.genesis_txs[1].hash})
 
     def test_vertices_count(self):
-        _set_test_mode(TestMode.TEST_ALL_WEIGHT)
+        self.manager.daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
 
         blocks_count = 1
         txs_count = 2
@@ -532,7 +532,7 @@ def _add_new_block(self, parents=None):
         return block
 
     def test_best_block_tips_cache(self):
-        _set_test_mode(TestMode.TEST_ALL_WEIGHT)
+        self.manager.daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
         self.manager.wallet.unlock(b'MYPASS')
         spent_blocks = add_new_blocks(self.manager, 10)
         self.assertEqual(self.tx_storage._best_block_tips_cache, [spent_blocks[-1].hash])
@@ -544,7 +544,7 @@ def test_best_block_tips_cache(self):
         self.assertEqual(self.tx_storage._best_block_tips_cache, [latest_blocks[-1].hash])
 
     def test_topological_sort(self):
-        _set_test_mode(TestMode.TEST_ALL_WEIGHT)
+        self.manager.daa.TEST_MODE = TestMode.TEST_ALL_WEIGHT
         _total = 0
         blocks = add_new_blocks(self.manager, 1, advance_clock=1)
         _total += len(blocks)