IgnoreUpdatesContext: add on_first_enter and matching tests

This commit is contained in:
Laurent Monin
2024-05-24 13:06:57 +02:00
parent 5e499e5dcd
commit ef14f46744
2 changed files with 19 additions and 1 deletions

View File

@@ -561,15 +561,18 @@ class IgnoreUpdatesContext:
updates if it is `False`.
"""
def __init__(self, on_exit=None, on_enter=None):
def __init__(self, on_exit=None, on_enter=None, on_first_enter=None):
self._entered = 0
self._on_exit = on_exit
self._on_enter = on_enter
self._on_first_enter = on_first_enter
def __enter__(self):
self._entered += 1
if self._on_enter:
self._on_enter()
if self._entered == 1 and self._on_first_enter:
self._on_first_enter()
def __exit__(self, type, value, tb):
self._entered -= 1

View File

@@ -954,6 +954,21 @@ class IgnoreUpdatesContextTest(PicardTestCase):
with context:
self.assertEqual(len(on_enter.mock_calls), 2)
def test_run_on_first_enter(self):
on_first_enter = Mock()
context = IgnoreUpdatesContext(on_first_enter=on_first_enter)
with context:
on_first_enter.assert_called()
on_first_enter.assert_called_once_with()
def test_run_on_first_enter_nested(self):
on_first_enter = Mock()
context = IgnoreUpdatesContext(on_first_enter=on_first_enter)
with context:
on_first_enter.assert_called_once_with()
with context:
on_first_enter.assert_called_once_with()
def test_nested_with(self):
context = IgnoreUpdatesContext()
with context: