sjyangkevin commented on code in PR #74542:
URL: https://github.com/apache/airflow/pull/74542#discussion_r4254410113
##########
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:
Review Comment:
Thanks for pointing this out. The tests under this class have been cleanup,
re-wrote, and moved under `TestTemplater`.
##########
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:
After the above clean up, there are now two newly added test functions under
`TestTemplater`.
* test_loads_template_from_zip
* test_loads_templates_from_zip_and_directory_search_paths
`test_loads_template_from_zip` invoke `create_template_env` which will go to
the branch for building the zip loader in `_build_template_loader`. Then, it
checks for 4 path patterns that the query can be rendered.
`test_loads_templates_from_zip_and_directory_search_paths` check two
scenarios. A template in the zip is found, and a template in the folder is
found, when both are specifying in `searchpath`.
##########
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
Review Comment:
the test had been cleanup.
##########
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]
Review Comment:
thanks for pointing this out. The test had been cleanup.
--
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]
