From eca45fb8213c281addcb0f68b97ce37925960c0b Mon Sep 17 00:00:00 2001
From: Tony Meyer <tony.meyer@canonical.com>
Date: Thu, 12 Dec 2024 09:47:47 +1300
Subject: [PATCH] Go back to issuing a real warning.

---
 ops/main.py                  | 46 +++++++++++++++++++++++-------------
 test/test_main_invocation.py | 26 ++++++++++----------
 2 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/ops/main.py b/ops/main.py
index 6ee36c9ef..24978e19c 100644
--- a/ops/main.py
+++ b/ops/main.py
@@ -15,8 +15,9 @@
 """Support legacy ops.main.main() import."""
 
 import inspect
-import logging
-from typing import Any, Optional, Type
+import os
+import warnings
+from typing import Any, Optional, Type, Union
 
 import ops.charm
 
@@ -27,10 +28,9 @@
     CHARM_STATE_FILE,  # type: ignore[reportUnusedImport]
     _Dispatcher,  # type: ignore[reportUnusedImport]
     _get_event_args,  # type: ignore[reportUnusedImport]
+    logger,  # type: ignore[reportUnusedImport]
 )
 
-logger = logging.getLogger(__name__)
-
 
 def _top_frame():
     frame = inspect.currentframe()
@@ -56,13 +56,7 @@ def main(charm_class: Type[ops.charm.CharmBase], use_juju_for_storage: Optional[
     # This means that we need to delay emitting the warning until the framework
     # has been set up, so we wrap the charm and do it on instantiation. However,
     # this means that a regular warning call won't provide the correct filename
-    # and line number - we could wrap this in a custom formatter, but it's
-    # simpler to just use a logging.warning() call instead.
-    # Note that unit tests almost never call main(), so wouldn't show this
-    # warning, so we're relying on integration tests, which aren't going to be
-    # treating warning calls in any special way. That means that we can just
-    # output the warning directly as a log message, and the effect will be the
-    # same.
+    # and line number.
     # Note also that this will be logged with every event. Our assumption is
     # that this will be noticeable enough during integration testing that it
     # will get fixed before going into production.
@@ -72,12 +66,30 @@ def main(charm_class: Type[ops.charm.CharmBase], use_juju_for_storage: Optional[
     class DeprecatedMainCharmBase(charm_class):
         def __init__(self, *args: Any, **kwargs: Any):
             super().__init__(*args, **kwargs)
-            logger.warning(
-                '%s:%s: DeprecationWarning: Calling `ops.main()` is deprecated, '
-                'call `ops.main()` instead',
-                frame.f_code.co_filename,
-                frame.f_lineno,
-            )
+
+            _original_format = warnings.formatwarning
+
+            def custom_warning_formatter(
+                message: Union[str, Warning],
+                category: Type[Warning],
+                *args: Any,
+                **kwargs: Any,
+            ) -> str:
+                """Like the default formatter, but patch in the filename and line number."""
+                return (
+                    f'{frame.f_code.co_filename}:{frame.f_lineno}: '
+                    f'{category.__name__}: {message}'
+                )
+
+            try:
+                warnings.formatwarning = custom_warning_formatter
+                warnings.warn(
+                    'Calling `ops.main()` is deprecated, call `ops.main()` instead',
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            finally:
+                warnings.formatwarning = _original_format
 
     return _main.main(
         charm_class=DeprecatedMainCharmBase, use_juju_for_storage=use_juju_for_storage
diff --git a/test/test_main_invocation.py b/test/test_main_invocation.py
index 2e5d7f82f..4105b3c17 100644
--- a/test/test_main_invocation.py
+++ b/test/test_main_invocation.py
@@ -49,11 +49,11 @@ def test_top_level_import(charm_env: None):
         ops.main()  # type: ignore
 
 
-def test_top_level_import_legacy_call(charm_env: None, caplog: pytest.LogCaptureFixture):
+def test_top_level_import_legacy_call(charm_env: None):
     import ops
 
-    ops.main.main(ops.CharmBase)
-    assert 'DeprecationWarning: Calling `ops.main()` is deprecated' in caplog.records[0].message
+    with pytest.deprecated_call():
+        ops.main.main(ops.CharmBase)
 
     with pytest.raises(TypeError):
         ops.main.main()  # type: ignore
@@ -68,11 +68,11 @@ def test_submodule_import(charm_env: None):
         ops.main()  # type: ignore
 
 
-def test_submodule_import_legacy_call(charm_env: None, caplog: pytest.LogCaptureFixture):
+def test_submodule_import_legacy_call(charm_env: None):
     import ops.main
 
-    ops.main.main(ops.CharmBase)
-    assert 'DeprecationWarning: Calling `ops.main()` is deprecated' in caplog.records[0].message
+    with pytest.deprecated_call():
+        ops.main.main(ops.CharmBase)
 
     with pytest.raises(TypeError):
         ops.main.main()  # type: ignore
@@ -87,23 +87,21 @@ def test_import_from_top_level_module(charm_env: None):
         main()  # type: ignore
 
 
-def test_import_from_top_level_module_legacy_call(
-    charm_env: None, caplog: pytest.LogCaptureFixture
-):
+def test_import_from_top_level_module_legacy_call(charm_env: None):
     from ops import main
 
-    main.main(ops.CharmBase)
-    assert 'DeprecationWarning: Calling `ops.main()` is deprecated' in caplog.records[0].message
+    with pytest.deprecated_call():
+        main.main(ops.CharmBase)
 
     with pytest.raises(TypeError):
         main.main()  # type: ignore
 
 
-def test_legacy_import_from_submodule(charm_env: None, caplog: pytest.LogCaptureFixture):
+def test_legacy_import_from_submodule(charm_env: None):
     from ops.main import main
 
-    main(ops.CharmBase)
-    assert 'DeprecationWarning: Calling `ops.main()` is deprecated' in caplog.records[0].message
+    with pytest.deprecated_call():
+        main(ops.CharmBase)
 
     with pytest.raises(TypeError):
         main()  # type: ignore