From 7caa6cfe312621dae341245d973c4fcda6920c0a Mon Sep 17 00:00:00 2001
From: Matthias <xmatthias@outlook.com>
Date: Sat, 17 Oct 2020 11:40:01 +0200
Subject: [PATCH] Add tests for pairlock

---
 freqtrade/persistence/models.py | 16 ++++++++++--
 freqtrade/strategy/interface.py | 10 +------
 tests/test_persistence.py       | 46 ++++++++++++++++++++++++++++++++-
 3 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/freqtrade/persistence/models.py b/freqtrade/persistence/models.py
index b2f8f4274..4394b783a 100644
--- a/freqtrade/persistence/models.py
+++ b/freqtrade/persistence/models.py
@@ -678,11 +678,23 @@ class PairLock(_DECL_BASE):
     active = Column(Boolean, nullable=False, default=True)
 
     def __repr__(self):
-        lock_time = self.open_date.strftime(DATETIME_PRINT_FORMAT)
-        lock_end_time = self.open_date.strftime(DATETIME_PRINT_FORMAT)
+        lock_time = self.lock_time.strftime(DATETIME_PRINT_FORMAT)
+        lock_end_time = self.lock_end_time.strftime(DATETIME_PRINT_FORMAT)
         return (f'PairLock(id={self.id}, pair={self.pair}, lock_time={lock_time}, '
                 f'lock_end_time={lock_end_time})')
 
+    @staticmethod
+    def lock_pair(pair: str, until: datetime, reason: str = None) -> None:
+        lock = PairLock(
+            pair=pair,
+            lock_time=datetime.now(timezone.utc),
+            lock_end_time=until,
+            reason=reason,
+            active=True
+        )
+        PairLock.session.add(lock)
+        PairLock.session.flush()
+
     @staticmethod
     def get_pair_locks(pair: str, now: Optional[datetime] = None) -> List['PairLock']:
         """
diff --git a/freqtrade/strategy/interface.py b/freqtrade/strategy/interface.py
index d9485e27a..36abfd05a 100644
--- a/freqtrade/strategy/interface.py
+++ b/freqtrade/strategy/interface.py
@@ -287,15 +287,7 @@ class IStrategy(ABC):
         :param until: datetime in UTC until the pair should be blocked from opening new trades.
                 Needs to be timezone aware `datetime.now(timezone.utc)`
         """
-        lock = PairLock(
-            pair=pair,
-            lock_time=datetime.now(timezone.utc),
-            lock_end_time=until,
-            reason=reason,
-            active=True
-        )
-        PairLock.session.add(lock)
-        PairLock.session.flush()
+        PairLock.lock_pair(pair, until, reason)
 
     def unlock_pair(self, pair: str) -> None:
         """
diff --git a/tests/test_persistence.py b/tests/test_persistence.py
index 4216565ac..6ac1e36a4 100644
--- a/tests/test_persistence.py
+++ b/tests/test_persistence.py
@@ -1,5 +1,6 @@
 # pragma pylint: disable=missing-docstring, C0103
 import logging
+from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock
 
 import arrow
@@ -8,7 +9,7 @@ from sqlalchemy import create_engine
 
 from freqtrade import constants
 from freqtrade.exceptions import DependencyException, OperationalException
-from freqtrade.persistence import Order, Trade, clean_dry_run_db, init_db
+from freqtrade.persistence import Order, PairLock, Trade, clean_dry_run_db, init_db
 from tests.conftest import create_mock_trades, log_has, log_has_re
 
 
@@ -1158,3 +1159,46 @@ def test_select_order(fee):
     assert order.ft_order_side == 'stoploss'
     order = trades[4].select_order('sell', False)
     assert order is None
+
+
+@pytest.mark.usefixtures("init_persistence")
+def test_PairLock(default_conf):
+    # No lock should be present
+    assert len(PairLock.query.all()) == 0
+
+    pair = 'ETH/BTC'
+    assert not PairLock.is_pair_locked(pair)
+    PairLock.lock_pair(pair, arrow.utcnow().shift(minutes=4).datetime)
+    # ETH/BTC locked for 4 minutes
+    assert PairLock.is_pair_locked(pair)
+
+    # XRP/BTC should not be locked now
+    pair = 'XRP/BTC'
+    assert not PairLock.is_pair_locked(pair)
+
+    # Unlocking a pair that's not locked should not raise an error
+    PairLock.unlock_pair(pair)
+
+    # Unlock original pair
+    pair = 'ETH/BTC'
+    PairLock.unlock_pair(pair)
+    assert not PairLock.is_pair_locked(pair)
+
+    pair = 'BTC/USDT'
+    # Lock until 14:30
+    lock_time = datetime(2020, 5, 1, 14, 30, 0, tzinfo=timezone.utc)
+    PairLock.lock_pair(pair, lock_time)
+    # Lock is in the past, so we must fake the lock
+    lock = PairLock.query.filter(PairLock.pair == pair).first()
+    lock.lock_time = lock_time - timedelta(hours=2)
+
+    assert not PairLock.is_pair_locked(pair)
+    assert PairLock.is_pair_locked(pair, lock_time + timedelta(minutes=-10))
+    assert PairLock.is_pair_locked(pair, lock_time + timedelta(minutes=-50))
+
+    # Should not be locked after time expired
+    assert not PairLock.is_pair_locked(pair, lock_time + timedelta(minutes=10))
+
+    locks = PairLock.get_pair_locks(pair, lock_time + timedelta(minutes=-2))
+    assert len(locks) == 1
+    assert 'PairLock' in str(locks[0])