diff --git a/freqtrade/persistence/__init__.py b/freqtrade/persistence/__init__.py index 595090540..6205174a7 100644 --- a/freqtrade/persistence/__init__.py +++ b/freqtrade/persistence/__init__.py @@ -4,4 +4,5 @@ from freqtrade.persistence.key_value_store import KeyStoreKeys, KeyValueStore from freqtrade.persistence.models import init_db from freqtrade.persistence.pairlock_middleware import PairLocks from freqtrade.persistence.trade_model import LocalTrade, Order, Trade -from freqtrade.persistence.usedb_context import disable_database_use, enable_database_use +from freqtrade.persistence.usedb_context import (FtNoDBContext, disable_database_use, + enable_database_use) diff --git a/freqtrade/persistence/usedb_context.py b/freqtrade/persistence/usedb_context.py index 4d8c9e539..6fffd2fb5 100644 --- a/freqtrade/persistence/usedb_context.py +++ b/freqtrade/persistence/usedb_context.py @@ -20,3 +20,14 @@ def enable_database_use() -> None: PairLocks.use_db = True PairLocks.timeframe = '' Trade.use_db = True + + +class FtNoDBContext: + def __init__(self, timeframe: str = ''): + self.timeframe = timeframe + + def __enter__(self): + disable_database_use(self.timeframe) + + def __exit__(self, exc_type, exc_val, exc_tb): + enable_database_use() diff --git a/tests/persistence/test_db_context.py b/tests/persistence/test_db_context.py new file mode 100644 index 000000000..5145268f9 --- /dev/null +++ b/tests/persistence/test_db_context.py @@ -0,0 +1,24 @@ +import pytest + +from freqtrade.persistence import FtNoDBContext, PairLocks, Trade + + +@pytest.mark.parametrize('timeframe', ['', '5m', '1d']) +def test_FtNoDBContext(timeframe): + assert Trade.use_db is True + assert PairLocks.use_db is True + assert PairLocks.timeframe == '' + + with FtNoDBContext(timeframe): + assert Trade.use_db is False + assert PairLocks.use_db is False + assert PairLocks.timeframe == timeframe + + with FtNoDBContext(): + assert Trade.use_db is False + assert PairLocks.use_db is False + assert PairLocks.timeframe == '' + + assert Trade.use_db is True + assert PairLocks.use_db is True +