support snapshot updates

Author: khvn26

src/common/test_tools/plugin.py

@@ -5,7 +5,16 @@
 from prometheus_client.metrics import MetricWrapperBase
 from pyfakefs.fake_filesystem import FakeFilesystem
 
-from common.test_tools.types import AssertMetricFixture, SnapshotFixture
+from common.test_tools.types import AssertMetricFixture, Snapshot, SnapshotFixture
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    group = parser.getgroup("snapshot")
+    group.addoption(
+        "--snapshot-update",
+        action="store_true",
+        help="Update snapshot files instead of testing against them.",
+    )
 
 
 def assert_metric_impl() -> Generator[AssertMetricFixture, None, None]:
@@ -72,10 +81,11 @@ def flagsmith_markers_marked(
 
 @pytest.fixture
 def snapshot(request: pytest.FixtureRequest) -> SnapshotFixture:
-    def _get_snapshot(name: str = "") -> str:
+    for_update = request.config.getoption("--snapshot-update")
+
+    def _get_snapshot(name: str = "") -> Snapshot:
         snapshot_name = name or f"{request.node.name}.txt"
-        return open(
-            request.path.parent / f"snapshots/{snapshot_name}",
-        ).read()
+        snapshot_path = request.path.parent / f"snapshots/{snapshot_name}"
+        return Snapshot(snapshot_path, for_update=for_update)
 
     return _get_snapshot

src/common/test_tools/types.py

@@ -1,5 +1,8 @@
+from pathlib import Path
 from typing import Protocol
 
+import pytest
+
 
 class AssertMetricFixture(Protocol):
     def __call__(
@@ -12,4 +15,24 @@ def __call__(
 
 
 class SnapshotFixture(Protocol):
-    def __call__(self, name: str = "") -> str: ...
+    def __call__(self, name: str = "") -> "Snapshot": ...
+
+
+class Snapshot:
+    def __init__(self, path: Path, for_update: bool) -> None:
+        self.path = path
+        self.content = open(path).read()
+        self.for_update = for_update
+
+    def __eq__(self, other: object) -> bool:
+        eq = self.content == other
+        if not eq and isinstance(other, str) and self.for_update:
+            with open(self.path, "w") as f:
+                f.write(other)
+                pytest.xfail(reason=f"Snapshot updated: {self.path}")
+        return eq
+
+    def __str__(self) -> str:
+        return self.content
+
+    __repr__ = __str__