Add form data to query context generator

tests/common/query_context_generator.py

@@ -16,7 +16,7 @@
 # under the License.
 import copy
 import dataclasses
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from superset.common.chart_data import ChartDataResultType
 from superset.utils.core import AnnotationType, DTTM_ALIAS, TimeRangeEndpoint
@@ -242,7 +242,9 @@ def generate(
         add_time_offsets: bool = False,
         table_id=1,
         table_type="table",
+        form_data: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
+        form_data = form_data or {}
         table_name = query_name.split(":")[0]
         table = self.get_table(table_name, table_id, table_type)
         return {
@@ -253,7 +255,8 @@ def generate(
                 )
             ],
             "result_type": ChartDataResultType.FULL,
+            "form_data": form_data,
         }
 
-    def get_table(self, name, id, type):
-        return Table(id, type, name)
+    def get_table(self, name, id_, type_):
+        return Table(id_, type_, name)

tests/integration_tests/fixtures/query_context.py

@@ -14,21 +14,22 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from tests.common.query_context_generator import QueryContextGenerator
 from tests.integration_tests.base_tests import SupersetTestCase
 
 
 class QueryContextGeneratorInteg(QueryContextGenerator):
-    def get_table(self, name, id, type):
+    def get_table(self, name, id_, type_):
         return SupersetTestCase.get_table(name=name)
 
 
 def get_query_context(
     query_name: str,
     add_postprocessing_operations: bool = False,
     add_time_offsets: bool = False,
+    form_data: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, Any]:
     """
     Create a request payload for retrieving a QueryContext object via the
@@ -40,8 +41,12 @@ def get_query_context(
     :param datasource_type: type of datasource to query.
     :param add_postprocessing_operations: Add post-processing operations to QueryObject
     :param add_time_offsets: Add time offsets to QueryObject(advanced analytics)
+    :param form_data: chart metadata
     :return: Request payload
     """
     return QueryContextGeneratorInteg().generate(
-        query_name, add_postprocessing_operations, add_time_offsets
+        query_name=query_name,
+        add_postprocessing_operations=add_postprocessing_operations,
+        add_time_offsets=add_time_offsets,
+        form_data=form_data,
     )