Skip to content

Commit

Permalink
add tests for rendered dataframe
Browse files Browse the repository at this point in the history
  • Loading branch information
ephraimbuddy committed Mar 24, 2024
1 parent 8d72e61 commit 987d980
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions tests/models/test_renderedtifields.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@

from airflow import settings
from airflow.configuration import conf
from airflow.decorators import task as task_decorator
from airflow.models import Variable
from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
from airflow.operators.bash import BashOperator
Expand Down Expand Up @@ -190,6 +191,28 @@ def test_secrets_are_masked_when_large_string(self, dag_maker):
assert ti.run_id == rtif.run_id
assert "***" in rtif.rendered_fields.get("bash_command")

@mock.patch("airflow.models.BaseOperator.render_template")
def test_pandas_dataframes_works_with_the_string_compare(self, render_mock, dag_maker):
import pandas

render_mock.return_value = pandas.DataFrame({"a": [1, 2, 3]})
with dag_maker("test_serialized_rendered_fields"):

@task_decorator
def generate_pd():
return pandas.DataFrame({"a": [1, 2, 3]})

@task_decorator
def consume_pd(data):
return data

consume_pd(generate_pd())

dr = dag_maker.create_dagrun()
ti, ti2 = dr.task_instances
rtif = RTIF(ti=ti2)
rtif.write()

@pytest.mark.parametrize(
"rtif_num, num_to_keep, remaining_rtifs, expected_query_count",
[
Expand Down

0 comments on commit 987d980

Please sign in to comment.