improve tests and decorator

Author: eschutho

superset/utils/decorators.py

@@ -62,7 +62,10 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
     return decorate
 
 
-def logs_context(**ctx_kwargs: int | str | UUID | None) -> Callable[..., Any]:
+def logs_context(
+    context_func: Callable[..., dict[Any, Any]] | None = None,
+    **ctx_kwargs: int | str | UUID | None,
+) -> Callable[..., Any]:
     """
     Takes arguments and adds them to the global logs_context.
     This is for logging purposes only and values should not be relied on or mutated
@@ -72,35 +75,60 @@ def decorate(f: Callable[..., Any]) -> Callable[..., Any]:
         def wrapped(*args: Any, **kwargs: Any) -> Any:
             if not hasattr(g, "logs_context"):
                 g.logs_context = {}
+
+            # limit data that can be saved to logs_context
+            # in order to prevent antipatterns
             available_logs_context_keys = [
                 "slice_id",
                 "dashboard_id",
+                "dataset_id",
                 "execution_id",
                 "report_schedule_id",
             ]
+            # set value from kwargs from
+            # wrapper function if it exists
+            # e.g. @logs_context()
+            #      def my_func(slice_id=None, **kwargs)
+            #
+            #      my_func(slice_id=2)
             logs_context_data = {
                 key: val
                 for key, val in kwargs.items()
                 if key in available_logs_context_keys
+                if val is not None
             }
 
-            # if keys are passed in to decorator directly, add them to logs_context
-            # by overriding values from kwargs
-            # e.g. @logs_context(slice_id=1, dashboard_id=1)
-            for key in available_logs_context_keys:
-                if ctx_kwargs.get(key) is not None:
-                    logs_context_data[key] = ctx_kwargs[key]
-                    try:
-                        # override value from local kwargs from
-                        # function execution if it exists
-                        # e.g. @logs_context(slice_id=1, dashboard_id=1)
-                        #      def my_func(slice_id=None, **kwargs)
-                        #
-                        #      my_func(slice_id=2)
-                        logs_context_data[key] = locals()["kwargs"][key]
-                    except KeyError:
-                        # do nothing if the key doesn't exist
-                        pass
+            try:
+                # if keys are passed in to decorator directly, add them to logs_context
+                # by overriding values from kwargs
+                # e.g. @logs_context(slice_id=1, dashboard_id=1)
+                logs_context_data.update(
+                    {
+                        key: ctx_kwargs.get(key)
+                        for key in available_logs_context_keys
+                        if ctx_kwargs.get(key) is not None
+                    }
+                )
+
+                if context_func is not None:
+                    # if a context function is passed in, call it and add the
+                    # returned values to logs_context
+                    # context_func=lambda *args, **kwargs: {
+                    # "slice_id": 1, "dashboard_id": 1
+                    # }
+                    logs_context_data.update(
+                        {
+                            key: value
+                            for key, value in context_func(*args, **kwargs).items()
+                            if key in available_logs_context_keys
+                            if value is not None
+                        }
+                    )
+
+            except (TypeError, KeyError):
+                # do nothing if the key doesn't exist
+                # or context is not callable
+                pass
 
             g.logs_context.update(logs_context_data)
             return f(*args, **kwargs)

tests/unit_tests/notifications/slack_tests.py

@@ -21,9 +21,11 @@
 from flask import g
 
 
+@patch("superset.reports.notifications.slack.g")
 @patch("superset.reports.notifications.slack.logger")
 def test_send_slack(
     logger_mock: MagicMock,
+    flask_global_mock: MagicMock,
 ) -> None:
     # `superset.models.helpers`, a dependency of following imports,
     # requires app context
@@ -32,7 +34,7 @@ def test_send_slack(
     from superset.reports.notifications.slack import SlackNotification, WebClient
 
     execution_id = uuid.uuid4()
-    g.logs_context = {"execution_id": execution_id}
+    flask_global_mock.logs_context = {"execution_id": execution_id}
     content = NotificationContent(
         name="test alert",
         header_data={
@@ -82,6 +84,3 @@ def test_send_slack(
 ```
 """,
         )
-
-    # reset g.logs_context
-    g.logs_context = None

tests/unit_tests/utils/test_decorators.py

@@ -23,7 +23,6 @@
 from unittest.mock import call, Mock, patch
 
 import pytest
-from flask import g
 
 from superset import app
 from superset.utils import decorators
@@ -89,7 +88,10 @@ def my_func(response: ResponseValues, *args: Any, **kwargs: Any) -> str:
             mock.assert_called_once_with(expected_result, 1)
 
 
-def test_context_decorator() -> None:
+@patch("superset.utils.decorators.g")
+def test_context_decorator(flask_g_mock) -> None:
+    flask_g_mock.logs_context = {}
+
     @decorators.logs_context()
     def myfunc(*args, **kwargs) -> str:
         return "test"
@@ -99,46 +101,85 @@ def myfunc(*args, **kwargs) -> str:
     def myfunc_with_kwargs(*args, **kwargs) -> str:
         return "test"
 
+    @decorators.logs_context(bad_context=1)
+    def myfunc_with_dissallowed_kwargs(*args, **kwargs) -> str:
+        return "test"
+
+    @decorators.logs_context(
+        context_func=lambda *args, **kwargs: {"slice_id": kwargs["chart_id"]}
+    )
+    def myfunc_with_context(*args, **kwargs) -> str:
+        return "test"
+
     # should not add any data to the global.logs_context scope
     myfunc(1, 1)
-    assert g.logs_context == {}
+    assert flask_g_mock.logs_context == {}
 
     # should add dashboard_id to the global.logs_context scope
     myfunc(1, 1, dashboard_id=1)
-    assert g.logs_context == {"dashboard_id": 1}
-    g.logs_context = {}
+    assert flask_g_mock.logs_context == {"dashboard_id": 1}
+    # reset logs_context
+    flask_g_mock.logs_context = {}
 
     # should add slice_id to the global.logs_context scope
     myfunc(1, 1, slice_id=1)
-    assert g.logs_context == {"slice_id": 1}
-    g.logs_context = {}
+    assert flask_g_mock.logs_context == {"slice_id": 1}
+    # reset logs_context
+    flask_g_mock.logs_context = {}
 
     # should add execution_id to the global.logs_context scope
     myfunc(1, 1, execution_id=1)
-    assert g.logs_context == {"execution_id": 1}
-    g.logs_context = {}
+    assert flask_g_mock.logs_context == {"execution_id": 1}
+    # reset logs_context
+    flask_g_mock.logs_context = {}
 
     # should add all three to the global.logs_context scope
     myfunc(1, 1, dashboard_id=1, slice_id=1, execution_id=1)
-    assert g.logs_context == {"dashboard_id": 1, "slice_id": 1, "execution_id": 1}
-    g.logs_context = {}
+    assert flask_g_mock.logs_context == {
+        "dashboard_id": 1,
+        "slice_id": 1,
+        "execution_id": 1,
+    }
+    # reset logs_context
+    flask_g_mock.logs_context = {}
 
     # should overwrite existing values in the global.logs_context scope
-    g.logs_context = {"dashboard_id": 2, "slice_id": 2, "execution_id": 2}
+    flask_g_mock.logs_context = {"dashboard_id": 2, "slice_id": 2, "execution_id": 2}
     myfunc(1, 1, dashboard_id=3, slice_id=3, execution_id=3)
-    assert g.logs_context == {"dashboard_id": 3, "slice_id": 3, "execution_id": 3}
-    g.logs_context = {}
+    assert flask_g_mock.logs_context == {
+        "dashboard_id": 3,
+        "slice_id": 3,
+        "execution_id": 3,
+    }
+    # reset logs_context
+    flask_g_mock.logs_context = {}
+
+    # should not add a value that does not exist in the global.logs_context scope
+    myfunc_with_dissallowed_kwargs()
+    assert flask_g_mock.logs_context == {}
+    # reset logs_context
+    flask_g_mock.logs_context = {}
 
     # should be able to add values to the decorator function directly
     myfunc_with_kwargs()
-    assert g.logs_context["dashboard_id"] == 1
-    assert g.logs_context["slice_id"] == 1
-    assert isinstance(g.logs_context["execution_id"], uuid.UUID)
-    g.logs_context = {}
+    assert flask_g_mock.logs_context["dashboard_id"] == 1
+    assert flask_g_mock.logs_context["slice_id"] == 1
+    assert isinstance(flask_g_mock.logs_context["execution_id"], uuid.UUID)
+    # reset logs_context
+    flask_g_mock.logs_context = {}
 
     # should be able to add values to the decorator function directly
-    # and overwrite with any kwargs passed in
+    # and it will overwrite any kwargs passed into the decorated function
     myfunc_with_kwargs(execution_id=4)
 
-    assert g.logs_context == {"dashboard_id": 1, "slice_id": 1, "execution_id": 4}
-    g.logs_context = {}
+    assert flask_g_mock.logs_context["dashboard_id"] == 1
+    assert flask_g_mock.logs_context["slice_id"] == 1
+    assert isinstance(flask_g_mock.logs_context["execution_id"], uuid.UUID)
+    # reset logs_context
+    flask_g_mock.logs_context = {}
+
+    # should be able to pass a callable context to the decorator
+    myfunc_with_context(chart_id=1)
+    assert flask_g_mock.logs_context == {"slice_id": 1}
+    # reset logs_context
+    flask_g_mock.logs_context = {}