Repository: apache/superset
File: tests/integration_tests/security/row_level_security_tests.py (Python)

Code before:

        db_tables = db.session.query(SqlaTable).all()

        db_table_names = {t.name for t in db_tables}
        received_tables = {table["text"] for table in result}

        assert data["count"] == len(db_tables)
        assert len(result) == len(db_tables)
        assert db_table_names == received_tables

    def test_rls_tables_related_api_with_filter(self):
        self.login(ADMIN_USERNAME)
        # Test with filter that should match 'birth_names'
        params = prison.dumps({"filter": "birth", "page": 0, "page_size": 100})
        rv = self.client.get(f"/api/v1/rowlevelsecurity/related/tables?q={params}")
        assert rv.status_code == 200
        data = json.loads(rv.data.decode("utf-8"))
        result = data["result"]
        received_tables = {table["text"] for table in result}
        # Should only return tables with 'birth' in the name
        assert all("birth" in table_name.lower() for table_name in received_tables)
        assert len(result) >= 1  # At least birth_names should be returned

        # Test with filter that should match nothing
        params = prison.dumps({"filter": "nonexistent", "page": 0, "page_size": 100})
        rv = self.client.get(f"/api/v1/rowlevelsecurity/related/tables?q={params}")
        assert rv.status_code == 200
        data = json.loads(rv.data.decode("utf-8"))
        result = data["result"]
        assert len(result) == 0
        assert data["count"] == 0

    def test_rls_roles_related_api(self):
        self.login(ADMIN_USERNAME)
        params = prison.dumps({"page": 0, "page_size": 100})

        rv = self.client.get(f"/api/v1/rowlevelsecurity/related/roles?q={params}")
        assert rv.status_code == 200
        data = json.loads(rv.data.decode("utf-8"))
        result = data["result"]

        db_role_names = {r.name for r in security_manager.get_all_roles()}
        received_roles = {role["text"] for role in result}

        assert data["count"] == len(db_role_names)
        assert len(result) == len(db_role_names)
        assert db_role_names == received_roles

    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
    @pytest.mark.usefixtures("load_energy_table_with_slice")
    @mock.patch(

Review comment (bug):
nit: I think I'd prefer breaking this up into a separate test such that if there is a failure that it's pretty clear which scenario broke instead of having to investigate if it's not returning the right item or empty is broken.

Code after:

        db_tables = db.session.query(SqlaTable).all()

        db_table_names = {t.name for t in db_tables}
        received_tables = {table["text"] for table in result}

        assert data["count"] == len(db_tables)
        assert len(result) == len(db_tables)
        assert db_table_names == received_tables

    def test_rls_tables_related_api_with_filter_matching_birth(self):
        self.login(ADMIN_USERNAME)
        # Test with filter that should match 'birth_names'
        params = prison.dumps({"filter": "birth", "page": 0, "page_size": 100})
        rv = self.client.get(f"/api/v1/rowlevelsecurity/related/tables?q={params}")
        assert rv.status_code == 200
        data = json.loads(rv.data.decode("utf-8"))
        result = data["result"]
        received_tables = {table["text"] for table in result}
        # Should only return tables with 'birth' in the name
        assert all("birth" in table_name.lower() for table_name in received_tables)
        assert len(result) >= 1  # At least birth_names should be returned

    def test_rls_tables_related_api_with_filter_no_matches(self):
        self.login(ADMIN_USERNAME)
        # Test with filter that should match nothing
        params = prison.dumps({"filter": "nonexistent", "page": 0, "page_size": 100})
        rv = self.client.get(f"/api/v1/rowlevelsecurity/related/tables?q={params}")
        assert rv.status_code == 200
        data = json.loads(rv.data.decode("utf-8"))
        result = data["result"]
        assert len(result) == 0
        assert data["count"] == 0

    def test_rls_roles_related_api(self):
        self.login(ADMIN_USERNAME)
        params = prison.dumps({"page": 0, "page_size": 100})

        rv = self.client.get(f"/api/v1/rowlevelsecurity/related/roles?q={params}")
        assert rv.status_code == 200
        data = json.loads(rv.data.decode("utf-8"))
        result = data["result"]

        db_role_names = {r.name for r in security_manager.get_all_roles()}
        received_roles = {role["text"] for role in result}

        assert data["count"] == len(db_role_names)
        assert len(result) == len(db_role_names)
        assert db_role_names == received_roles

    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")