From d6e8d0aad59bbe2bf1f46edb9abb30b7647f694f Mon Sep 17 00:00:00 2001
From: Tony Meyer <tony.meyer@canonical.com>
Date: Thu, 28 Mar 2024 16:30:00 +1300
Subject: [PATCH] Experiment with a tracking system (super ugly and not usable
 in practice right now, but good enough to run some tests).

---
 ops/framework.py |  4 ++++
 ops/testing.py   | 30 ++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/ops/framework.py b/ops/framework.py
index 07294ed7a..3d50aa727 100644
--- a/ops/framework.py
+++ b/ops/framework.py
@@ -958,6 +958,10 @@ def _reemit(self, single_event_path: Optional[str] = None):
             # scratch in the next path.
             self.framework._forget(event)  # type: ignore
 
+            # XXX Super hacky just for a PoC.
+            if isinstance(event, charm.HookEvent) and hasattr(observer, "_reset_attributes"):
+                observer._reset_attributes()
+
         if not deferred and last_event_path is not None:
             self._storage.drop_snapshot(last_event_path)
 
diff --git a/ops/testing.py b/ops/testing.py
index befd5d83a..baa7bb702 100644
--- a/ops/testing.py
+++ b/ops/testing.py
@@ -359,6 +359,36 @@ class TestEvents(self._charm_cls.on.__class__):
 
         class TestCharm(self._charm_cls):  # type: ignore
             on = TestEvents()
+            def __init__(self, *args: Any, **kwargs: Any):
+                super().__setattr__("_attributes", set(dir(self)))
+                super().__setattr__("_disabled", True)
+                super().__init__(*args, **kwargs)
+                super().__setattr__("_fresh_attributes", set(super().__getattribute__("_attributes")))
+                super().__setattr__("_disabled", False)
+
+            def __getattribute__(self, name: str) -> Any:
+                try:
+                    if super().__getattribute__("_disabled"):
+                        disabled = True
+                        attrs = {}
+                    else:
+                        attrs = super().__getattribute__("_attributes")
+                        disabled = False
+                except AttributeError:
+                    # Still setting up.
+                    pass
+                else:
+                    if not disabled and name not in attrs:
+                        raise AttributeError(f"You didn't set {name} in this event")
+                return super().__getattribute__(name)
+            
+            def __setattr__(self, name: str, value: Any) -> None:
+                if name != "_fresh_attributes":
+                    super().__getattribute__("_attributes").add(name)
+                return super().__setattr__(name, value)
+            
+            def _reset_attributes(self):
+                super().__setattr__("_attributes", set(super().__getattribute__("_fresh_attributes")))
 
         # Note: jam 2020-03-01 This is so that errors in testing say MyCharm has no attribute foo,
         # rather than TestCharm has no attribute foo.