ashb commented on code in PR #74542:
URL: https://github.com/apache/airflow/pull/74542#discussion_r4250963726
##########
task-sdk/tests/task_sdk/definitions/_internal/test_templater.py:
##########
@@ -300,6 +308,115 @@ def
test_do_render_template_fields_renders_list_values(self):
assert parent.items == ["first", "second"]
+def _write_zip(path: Path, members: dict[str, str]) -> Path:
+ with zipfile.ZipFile(path, "w") as zf:
+ for name, content in members.items():
+ zf.writestr(name, content)
+ return path
+
+
+def _write_template_dir(path: Path, templates: dict[str, str]) -> Path:
+ path.mkdir()
+ for name, content in templates.items():
+ (path / name).write_text(content)
+ return path
+
+
+class TestCreateTemplateEnvZip:
+ @pytest.mark.parametrize(
+ ("archive_name", "searchpath_parts", "template_name"),
+ [
+ pytest.param("dags.zip", (), "sql/query.sql", id="archive"),
+ pytest.param("dags.zip", ("sql",), "query.sql",
id="directory-in-archive"),
+ pytest.param("DAGS.ZIP", (), "sql/query.sql",
id="uppercase-extension"),
+ ],
+ )
+ def test_loads_template_from_zip(self, tmp_path, archive_name,
searchpath_parts, template_name):
+ archive = _write_zip(tmp_path / archive_name, {"sql/query.sql":
"SELECT {{ x }}"})
+
+ env = create_template_env(searchpath=[os.path.join(archive,
*searchpath_parts)])
+
+ assert env.get_template(template_name).render(x=1) == "SELECT 1"
+
+ @pytest.mark.parametrize(
+ "searchpath_parts",
+ [
+ pytest.param(("templates",), id="directory"),
+ pytest.param(("templates.zip",),
id="directory-named-like-archive"),
+ pytest.param(("not_a_zip.zip", "sql"), id="non-zip-file"),
+ pytest.param(("missing", "sql"), id="missing-path"),
+ ],
+ )
+ def test_keeps_filesystem_loader_without_zip(self, tmp_path,
searchpath_parts):
+ (tmp_path / "templates").mkdir()
+ (tmp_path / "templates.zip").mkdir()
+ (tmp_path / "not_a_zip.zip").write_text("plain text")
+ searchpath = os.path.join(tmp_path, *searchpath_parts)
+
+ env = create_template_env(searchpath=[searchpath])
+
+ assert type(env.loader) is jinja2.FileSystemLoader
+ assert env.loader.searchpath == [searchpath]
+
+ @pytest.mark.parametrize(("zip_first", "expected"), [(True, "from zip"),
(False, "from directory")])
+ def test_searches_zip_and_directory_in_order(self, tmp_path, zip_first,
expected):
+ archive = str(_write_zip(tmp_path / "dags.zip", {"query.sql": "from
zip"}))
+ directory = str(_write_template_dir(tmp_path / "templates",
{"query.sql": "from directory"}))
+
+ env = create_template_env(searchpath=[archive, directory] if zip_first
else [directory, archive])
+
+ assert env.get_template("query.sql").render() == expected
+
+ def test_template_not_found_lists_search_paths(self, tmp_path):
+ archive = str(_write_zip(tmp_path / "dags.zip", {"query.sql": "SELECT
1"}))
+ directory = str(_write_template_dir(tmp_path / "templates", {}))
+ env = create_template_env(searchpath=[archive, directory])
+
+ with pytest.raises(jinja2.TemplateNotFound) as ctx:
+ env.get_template("missing.sql")
+
+ assert ctx.value.name == "missing.sql"
+ assert ctx.value.message == f"'missing.sql' not found in search paths:
{archive!r}, {directory!r}"
Review Comment:
This is testing Jinja behaviour, not ours.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
