From 3ae80d3b98f9d27bc5a0a49bb7c87c3924281196 Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Tue, 8 Sep 2020 07:42:48 +0300
Subject: [PATCH] test: add rls jinja extra cache key test (#10810)

---
 tests/security_tests.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/security_tests.py b/tests/security_tests.py
index 64c6f32cc..50548fdeb 100644
--- a/tests/security_tests.py
+++ b/tests/security_tests.py
@@ -1020,7 +1020,7 @@ class TestRowLevelSecurity(SupersetTestCase):
             .filter(SqlaTable.table_name.in_(["energy_usage", "unicode_test"]))
             .all()
         )
-        self.rls_entry.clause = "value > 1"
+        self.rls_entry.clause = "value > {{ cache_key_wrapper(1) }}"
         self.rls_entry.roles.append(
             security_manager.find_role("Gamma")
         )  # db.session.query(Role).filter_by(name="Gamma").first())
@@ -1052,7 +1052,8 @@ class TestRowLevelSecurity(SupersetTestCase):
             extras={},
         )
         sql = tbl.get_query_str(query_obj)
-        self.assertIn("value > 1", sql)
+        assert tbl.get_extra_cache_keys(query_obj) == [1]
+        assert "value > 1" in sql
 
     def test_rls_filter_doesnt_alter_query(self):
         g.user = self.get_user(
@@ -1071,7 +1072,8 @@ class TestRowLevelSecurity(SupersetTestCase):
             extras={},
         )
         sql = tbl.get_query_str(query_obj)
-        self.assertNotIn("value > 1", sql)
+        assert tbl.get_extra_cache_keys(query_obj) == []
+        assert "value > 1" not in sql
 
     def test_multiple_table_filter_alters_another_tables_query(self):
         g.user = self.get_user(
@@ -1090,4 +1092,5 @@ class TestRowLevelSecurity(SupersetTestCase):
             extras={},
         )
         sql = tbl.get_query_str(query_obj)
-        self.assertIn("value > 1", sql)
+        assert tbl.get_extra_cache_keys(query_obj) == [1]
+        assert "value > 1" in sql