From f27155d9e57c621febda2237d2404f6d7fd86cb2 Mon Sep 17 00:00:00 2001
From: Ivan Shcheklein <shcheklein@gmail.com>
Date: Sun, 7 Aug 2022 19:18:29 -0700
Subject: [PATCH] fix logic for unlock on unlocked lock calls

---
 dvc/lock.py             | 11 +++++++++++
 tests/func/test_lock.py | 23 +++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/dvc/lock.py b/dvc/lock.py
index 3360001c04..17bbabc332 100644
--- a/dvc/lock.py
+++ b/dvc/lock.py
@@ -70,6 +70,8 @@ def lock(self):
         self._lock = True
 
     def unlock(self):
+        if not self.is_locked:
+            raise DvcException("Unlock called on an unlocked lock")
         self._lock = False
 
     @property
@@ -93,6 +95,7 @@ def __init__(self, lockfile, friendly=False, **kwargs):
         super().__init__(lockfile)
         self._friendly = friendly
         self._lock = None
+        self._lock_failed = False
 
     @property
     def files(self):
@@ -100,6 +103,7 @@ def files(self):
 
     def _do_lock(self):
         try:
+            self._lock_failed = False
             with Tqdm(
                 bar_format="{desc}",
                 disable=not self._friendly,
@@ -111,6 +115,7 @@ def _do_lock(self):
             ):
                 self._lock = zc.lockfile.LockFile(self._lockfile)
         except zc.lockfile.LockError:
+            self._lock_failed = True
             raise LockError(FAILED_TO_LOCK_MESSAGE)
 
     def lock(self):
@@ -120,6 +125,12 @@ def lock(self):
         lock_retry()
 
     def unlock(self):
+        if self._lock_failed:
+            assert self._lock is None
+            return
+
+        if not self.is_locked:
+            raise DvcException("Unlock called on an unlocked lock")
         self._lock.close()
         self._lock = None
 
diff --git a/tests/func/test_lock.py b/tests/func/test_lock.py
index b047404c06..461d2c23de 100644
--- a/tests/func/test_lock.py
+++ b/tests/func/test_lock.py
@@ -1,6 +1,7 @@
 import pytest
 
 from dvc.cli import main
+from dvc.exceptions import DvcException
 from dvc.lock import Lock, LockError
 
 
@@ -14,6 +15,28 @@ def test_with(tmp_dir, dvc, mocker):
             pass
 
 
+def test_(tmp_dir, dvc, mocker):
+    # patching to speedup tests
+    mocker.patch("dvc.lock.DEFAULT_TIMEOUT", 0.01)
+
+    lockfile = tmp_dir / dvc.tmp_dir / "lock"
+    lock = Lock(lockfile)
+    lock_ext = Lock(lockfile)
+
+    # It's a common scenario now to have lock unlocked and locked back (e.g. in
+    # repro of a stage) in with. We should see LockError exception here.
+    with pytest.raises(LockError), lock:
+        lock.unlock()
+        lock_ext.lock()  # imitate an exernal process has time to lock it
+        lock.lock()
+
+
+def test_unlock_unlocked_raises():
+    lock = Lock("lock")
+    with pytest.raises(DvcException):
+        lock.unlock()
+
+
 def test_cli(tmp_dir, dvc, mocker, caplog):
     # patching to speedup tests
     mocker.patch("dvc.lock.DEFAULT_TIMEOUT", 0.01)