fix tests

apache · Jun 11, 2024 · caaba6a · caaba6a
1 parent f0dc892
commit caaba6a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 19 deletions.
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
@@ -1377,7 +1377,7 @@ def values_for_column(
                 # automatically add a random alias to the projection because of the
                 # call to DISTINCT; others will uppercase the column names. This
                 # gives us a deterministic column name in the dataframe.
-                [target_col.get_sqla_col(template_processor=tp).label("column_values")]
+                target_col.get_sqla_col(template_processor=tp).label("column_values")
             )
             .select_from(tbl)
             .distinct()

diff --git a/tests/integration_tests/sqla_models_tests.py b/tests/integration_tests/sqla_models_tests.py
@@ -235,18 +235,21 @@ def test_jinja_metrics_and_calc_columns(self, mock_username):
             "'{{ 'xyz_' + time_grain }}' as time_grain",
             database=get_example_database(),
         )
-        TableColumn(
+        table_column = TableColumn(
             column_name="expr",
             expression="case when '{{ current_username() }}' = 'abc' "
             "then 'yes' else 'no' end",
             type="VARCHAR(100)",
             table=table,
         )
-        SqlMetric(
+        table_metric = SqlMetric(
             metric_name="count_timegrain",
             expression="count('{{ 'bar_' + time_grain }}')",
             table=table,
         )
+        db.session.add(table)
+        db.session.add(table_column)
+        db.session.add(table_metric)
         db.session.commit()
 
         sqla_query = table.get_sqla_query(**base_query_obj)
@@ -271,8 +274,9 @@ def test_jinja_metric_macro(self, mock_form_data_context):
         self.login(username="admin")
         table = self.get_table(name="birth_names")
         metric = SqlMetric(
-            metric_name="count_jinja_metric", expression="count(*)", table=table
+            metric_name="count_jinja_metric", expression="count(*)", table_id=table.id
         )
+        db.session.add(metric)
         db.session.commit()
 
         base_query_obj = {
@@ -347,7 +351,7 @@ def test_adhoc_metrics_and_calc_columns(self):
         with pytest.raises(QueryObjectValidationError):
             table.get_sqla_query(**base_query_obj)
         # Cleanup
-        db.session.delete(table)
+        db.session.query(SqlaTable).filter_by(table_name="test_validate_adhoc_sql").delete()
         db.session.commit()
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")

diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py
@@ -27,6 +27,7 @@
 import prison
 
 from freezegun import freeze_time
+from sqlalchemy import text
 from superset import db, security_manager
 from superset.connectors.sqla.models import SqlaTable  # noqa: F401
 from superset.db_engine_specs import BaseEngineSpec
@@ -212,20 +213,21 @@ def test_sql_json_cta_dynamic_db(self, ctas_method):
             db.session.commit()
             examples_db = get_example_database()
             with examples_db.get_sqla_engine() as engine:
-                data = engine.execute(
-                    f"SELECT * FROM admin_database.{tmp_table_name}"
-                ).fetchall()
-                names_count = engine.execute(
-                    f"SELECT COUNT(*) FROM birth_names"  # noqa: F541
-                ).first()
-                self.assertEqual(
-                    names_count[0], len(data)
-                )  # SQL_MAX_ROW not applied due to the SQLLAB_CTAS_NO_LIMIT set to True
-
-                # cleanup
-                engine.execute(f"DROP {ctas_method} admin_database.{tmp_table_name}")
-                examples_db.allow_ctas = old_allow_ctas
-                db.session.commit()
+                with engine.connect() as conn:
+                    data = conn.execute(
+                        text(f"SELECT * FROM admin_database.{tmp_table_name}")
+                    ).fetchall()
+                    names_count = conn.execute(
+                        text(f"SELECT COUNT(*) FROM birth_names")  # noqa: F541
+                    ).first()
+                    self.assertEqual(
+                        names_count[0], len(data)
+                    )  # SQL_MAX_ROW not applied due to the SQLLAB_CTAS_NO_LIMIT set to True
+
+                    # cleanup
+                    conn.execute(text(f"DROP {ctas_method} admin_database.{tmp_table_name}"))
+                    examples_db.allow_ctas = old_allow_ctas
+                    db.session.commit()
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_multi_sql(self):