feat: Migrates Pivot Table v1 to v2 (#23712)

tests/unit_tests/conftest.py

@@ -15,12 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=redefined-outer-name, import-outside-toplevel
-
+import functools
 import importlib
 import os
 import unittest.mock
 from collections.abc import Iterator
 from typing import Any, Callable
+from unittest.mock import patch
 
 import pytest
 from _pytest.fixtures import SubRequest
@@ -33,7 +34,7 @@ from superset import security_manager
 from superset.app import SupersetApp
 from superset.common.chart_data import ChartDataResultType
 from superset.common.query_object_factory import QueryObjectFactory
-from superset.extensions import appbuilder
+from superset.extensions import appbuilder, feature_flag_manager
 from superset.initialization import SupersetAppInitializer
 
 
@@ -164,3 +165,38 @@ def dummy_query_object(request, app_context):
         _datasource_dao=unittest.mock.Mock(),
         session_maker=unittest.mock.Mock(),
     ).create(parent_result_type=result_type, **query_object)
+
+
+def with_feature_flags(**mock_feature_flags):
+    """
+    Use this decorator to mock feature flags in tests.integration_tests.
+
+    Usage:
+
+        class TestYourFeature(SupersetTestCase):
+
+            @with_feature_flags(YOUR_FEATURE=True)
+            def test_your_feature_enabled(self):
+                self.assertEqual(is_feature_enabled("YOUR_FEATURE"), True)
+
+            @with_feature_flags(YOUR_FEATURE=False)
+            def test_your_feature_disabled(self):
+                self.assertEqual(is_feature_enabled("YOUR_FEATURE"), False)
+    """
+
+    def mock_get_feature_flags():
+        feature_flags = feature_flag_manager._feature_flags or {}
+        return {**feature_flags, **mock_feature_flags}
+
+    def decorate(test_fn):
+        def wrapper(*args, **kwargs):
+            with patch.object(
+                feature_flag_manager,
+                "get_feature_flags",
+                side_effect=mock_get_feature_flags,
+            ):
+                test_fn(*args, **kwargs)
+
+        return functools.update_wrapper(wrapper, test_fn)
+
+    return decorate