diff --git a/pytest_generator/generate_tests.py b/pytest_generator/generate_tests.py
index d17b560..72b8572 100644
--- a/pytest_generator/generate_tests.py
+++ b/pytest_generator/generate_tests.py
@@ -9,15 +9,16 @@
 from types import MethodType, ModuleType
 from typing import Any, Callable, Optional
 
-import jinja2
+from jinja2 import Environment, FileSystemLoader, select_autoescape
 
-env = jinja2.Environment(
-    loader=jinja2.FileSystemLoader(Path(__file__).parent.parent / "templates"),
-    autoescape=jinja2.select_autoescape(),
+env = Environment(
+    loader=FileSystemLoader(Path(__file__).parent.parent / "templates"),
+    autoescape=select_autoescape(),
 )
 
 
 def generate(root: Path):
+    """Generate tests for all modules under `root`."""
     if not root.is_absolute():
         root = root.absolute()
 
diff --git a/tests/pytest_generator/test_generate_tests.py b/tests/pytest_generator/test_generate_tests.py
index 4c34705..449e51d 100644
--- a/tests/pytest_generator/test_generate_tests.py
+++ b/tests/pytest_generator/test_generate_tests.py
@@ -1,6 +1,6 @@
 import pytest
-from pytest_generator import generate_tests
 
+from pytest_generator import generate_tests
 
 
 @pytest.mark.parametrize(("expected_output"), [("x")])
@@ -9,14 +9,12 @@ def test_generate_tests_generate(expected_output):
     assert return_value == expected_output
 
 
-
 @pytest.mark.parametrize(("expected_output"), [("x")])
 def test_generate_tests__get_local_members(expected_output):
     return_value = generate_tests._get_local_members()
     assert return_value == expected_output
 
 
-
 @pytest.mark.parametrize(("expected_output"), [("x")])
 def test_generate_tests__collect_members_and_write_tests(expected_output):
     return_value = generate_tests._collect_members_and_write_tests()
diff --git a/tests/pytest_generator/test_main.py b/tests/pytest_generator/test_main.py
index fef9e9c..e4a60f5 100644
--- a/tests/pytest_generator/test_main.py
+++ b/tests/pytest_generator/test_main.py
@@ -1,6 +1,6 @@
 import pytest
-from pytest_generator import main
 
+from pytest_generator import main
 
 
 @pytest.mark.parametrize(("expected_output"), [("x")])