def test_generate_role_session_name():
dag = DAG("a-test-dag")
task = airflow_docker.operator.Operator(
dag=dag,
task_id="some-task",
image="hello-world",
start_date=datetime.datetime(2019, 2, 14, 15),
)
ti = TaskInstance(task=task, execution_date=datetime.datetime(2019, 2, 14, 15))
dag_run = DagRun(dag_id=dag.dag_id)
dag_run.id = 5
context = {"dag": dag, "task_instance": ti, "dag_run": dag_run}
session_name = airflow_docker.ext.aws.role_assumption.generate_role_session_name(
context
)
assert "5__1__some-task" == session_name
def test_generate_role_session_name_long_task_id():
dag = DAG("a-test-dag")
task = airflow_docker.operator.Operator(
dag=dag,
task_id=
"some-task-id-that-is-very-long-way-past-the-64-character-limit-foo-bar-baz",
image="hello-world",
start_date=datetime.datetime(2019, 2, 14, 15),
)
ti = TaskInstance(task=task,
execution_date=datetime.datetime(2019, 2, 14, 15))
dag_run = DagRun(dag_id=dag.dag_id)
dag_run.id = 5
context = {"dag": dag, "task_instance": ti, "dag_run": dag_run}
session_name = airflow_docker.ext.aws.role_assumption.generate_role_session_name(
context)
assert ("5__1__some-task-id-that-is-very-long-way-past-the-64-character-l"
== session_name)
