Revisiting “Database Operators Have SQL” test

Back in my first post, https://learningtotest.com/2019/06/13/3-easy-airflow-tests/, the third test made sure that there was valid SQL attached to the task.

I’ve come to realize that a downside to that particular snippet is that the output doesn’t tell you the name of the task that failed, but returns output like this:

tests/test_my_dag.py::test_database_operators_have_sql[db_operator0] FAILED [ 18%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator1] FAILED [ 27%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator2] FAILED [ 36%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator3] FAILED [ 45%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator4] FAILED [ 54%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator5] FAILED [ 63%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator6] FAILED [ 72%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator7] FAILED [ 81%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator8] FAILED [ 90%]
tests/test_my_dag.py::test_database_operators_have_sql[db_operator9] FAILED [100%]

That’s not very helpful! We can do better… Let’s generate a list of ids. It requires zipping up the two lists generated from dag.tasks & dag.task_ids though.

@mark.parametrize(
    "db_operator",
    [x for x in dag.tasks if isinstance(x, PostgresOperator)],
    ids=[x for x, y in zip(dag.task_ids, dag.tasks) if isinstance(y, PostgresOperator)],
)
def test_database_operators_have_sql(db_operator):
    """For all PostgresOperator task(s), verify that sql attribute returns
    non-empty value.
    """
    assert db_operator.sql.startswith("INSERT") or db_operator.sql.startswith(
        "TRUNCATE"
    )

Which now returns much more helpful output, telling us exactly which task failed uder test:

tests/test_my_dag.py::test_database_operators_have_sql[extract_table01] FAILED [ 64%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table02] FAILED [ 68%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table03] FAILED [ 72%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table04] FAILED [ 76%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table05] FAILED [ 80%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table06] FAILED [ 84%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table07] FAILED [ 88%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table08] FAILED [ 92%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table09] FAILED [ 96%]
tests/test_my_dag.py::test_database_operators_have_sql[extract_table10] FAILED [100%]

Ah, that’s much better… (Well, now I clearly need to go write some SQL!)