Repository: charmed-kubernetes/pytest-operator (language: python)

Issue:
fast-forward status-update hooks context
I found it was a useful pattern to temporarily speed up the status-update hook firing rate.
This looks like a job for a context manager.

Proposal:

```python
async with ops_test.fast_forward('10s'):
   # do stuff which might take some time to reflect on the workload status
   assert ops_test.model.wait_for_idle(...)
```


Discussion / hints:

Fix patch:
diff --git a/docs/reference.md b/docs/reference.md
index a96b1da..c88c385 100644
--- a/docs/reference.md
+++ b/docs/reference.md
@@ -248,6 +248,23 @@ Runs a Juju CLI command.
 Useful for cases where python-libjuju sees things differently than the Juju CLI. Will set
 `JUJU_MODEL`, so manually passing in `-m model-name` is unnecessary.
 
+#### `async def fast_forward(self, fast_interval: str = "10s", slow_interval: Optional[str] = None)`
+
+Temporarily speed up update-status firing rate for the current model.
+Returns an async context manager that temporarily sets update-status firing rate to `fast_interval`.
+If provided, when the context exits the update-status firing rate will be set to `slow_interval`. Otherwise, it will be set to the previous value.
+
+* If `fast_interval` is provided, the update-status firing rate will be set to that value upon entering the context. Default is 10s.
+* If `slow_interval` is provided, after the context exits the update-status firing rate will be set to that value; otherwise, to the value it had before the context was entered.
+
+It is effectively a shortcut for:
+```python
+    await ops_test.model.set_config({"update-status-hook-interval": <fast-interval>})
+
+    # do something
+
+    await ops_test.model.set_config({"update-status-hook-interval": <slow-interval>})
+```
 
 #### `def abort(self, *args, **kwargs)`
 


Test patch:
diff --git a/pytest_operator/plugin.py b/pytest_operator/plugin.py
index 5113eb0..b892559 100644
--- a/pytest_operator/plugin.py
+++ b/pytest_operator/plugin.py
@@ -1307,3 +1307,29 @@ def render_charms(
             self.render_charm(charm_path, include, exclude, context, **kwcontext)
             for charm_path in charm_paths
         ]
+
+    @contextlib.asynccontextmanager
+    async def fast_forward(
+        self, fast_interval: str = "10s", slow_interval: Optional[str] = None
+    ):
+        """Temporarily speed up update-status firing rate for the current model.
+
+        Returns an async context manager that temporarily sets update-status
+        firing rate to `fast_interval`.
+        If provided, when the context exits the update-status firing rate will
+        be set to `slow_interval`. Otherwise, it will be set to the previous
+        value.
+        """
+        model = self.model
+        if not model:
+            raise RuntimeError("No model currently set.")
+
+        update_interval_key = "update-status-hook-interval"
+        if slow_interval:
+            interval_after = slow_interval
+        else:
+            interval_after = (await model.get_config())[update_interval_key]
+
+        await model.set_config({update_interval_key: fast_interval})
+        yield
+        await model.set_config({update_interval_key: interval_after})
diff --git a/tests/integration/test_pytest_operator.py b/tests/integration/test_pytest_operator.py
index 5b6cdbc..9650d55 100644
--- a/tests/integration/test_pytest_operator.py
+++ b/tests/integration/test_pytest_operator.py
@@ -4,6 +4,8 @@
 
 import pytest
 
+from pytest_operator.plugin import OpsTest
+
 log = logging.getLogger(__name__)
 
 
@@ -127,6 +129,30 @@ async def test_run(ops_test):
     assert await ops_test.run("/usr/bin/rev", stdin=stdin) == (0, revd, "")
 
 
+@pytest.mark.parametrize(
+    "fast_interval, slow_interval",
+    (
+        ("10s", None),
+        ("10s", "10m"),
+        ("42s", "42m"),
+        ("41m", "41s"),  # odd but... why not.
+        ("43s", None),
+    ),
+)
+async def test_fast_forward(ops_test: OpsTest, fast_interval, slow_interval):
+    async def _get_rate():
+        return (await ops_test.model.get_config())["update-status-hook-interval"].value
+
+    previous_rate = await _get_rate()
+
+    async with ops_test.fast_forward(
+        fast_interval=fast_interval, slow_interval=slow_interval
+    ):
+        assert await _get_rate() == fast_interval
+
+    assert await _get_rate() == slow_interval or previous_rate
+
+
 @pytest.mark.abort_on_fail(abort_on_xfail=True)
 def test_abort():
     pytest.xfail("abort")