test/python/test_db_sql_preprocessor.py

   1 """
   2 Tests for SQL preprocessing.
   3 """
   4 from pathlib import Path
   5
   6 import pytest
   7
   8 @pytest.fixture
   9 def sql_factory(tmp_path):
  10     def _mk_sql(sql_body):
  11         (tmp_path / 'test.sql').write_text("""
  12           CREATE OR REPLACE FUNCTION test() RETURNS TEXT
  13           AS $$
  14           BEGIN
  15             {}
  16           END;
  17           $$ LANGUAGE plpgsql IMMUTABLE;""".format(sql_body))
  18         return 'test.sql'
  19
  20     return _mk_sql
  21
  22 @pytest.mark.parametrize("expr,ret", [
  23     ("'a'", 'a'),
  24     ("'{{db.partitions|join}}'", '012'),
  25     ("{% if 'country_name' in db.tables %}'yes'{% else %}'no'{% endif %}", "yes"),
  26     ("{% if 'xxx' in db.tables %}'yes'{% else %}'no'{% endif %}", "no"),
  27     ])
  28 def test_load_file_simple(sql_preprocessor, sql_factory, temp_db_conn, temp_db_cursor, expr, ret):
  29     sqlfile = sql_factory("RETURN {};".format(expr))
  30
  31     sql_preprocessor.run_sql_file(temp_db_conn, sqlfile)
  32
  33     assert temp_db_cursor.scalar('SELECT test()') == ret
  34
  35
  36 def test_load_file_with_params(sql_preprocessor, sql_factory, temp_db_conn, temp_db_cursor):
  37     sqlfile = sql_factory("RETURN '{{ foo }} {{ bar }}';")
  38
  39     sql_preprocessor.run_sql_file(temp_db_conn, sqlfile, bar='XX', foo='ZZ')
  40
  41     assert temp_db_cursor.scalar('SELECT test()') == 'ZZ XX'