Flask-Vite integration refactor before upstreaming (#5434)

* chore: refactor Vite-Flask integration

* fix: tests

* feature: handle exception when extension is not initialized

* chore: tests cleanup

* chore: more refactoring

* chore: rename outdir env var for consistency

* fix: rework extension init logic

* fix: relative module path

* fix: linting

* fix: incorrect usage of __all__

* fix: comment

* fix: config types

* chore: refactor ProdViteIntegration

- move public methods at the top of the class for clarity
- add helper for checking if asset chunks exist in manifest
- add helpers for checking if asset files exist in filesystem

* fix: align extension name to - convention

* feature: support for more source file extensions

* chore: better typing

* feature: auto-inject Vite dev tools in dev mode

* fix: typing

* fix: linting

* feature: auto-injected dev tools test

Author: edisile

.env

@@ -1,6 +1,6 @@
 PORT=8004
 VITE_PORT=5004
-VITE_OUTPUT_DIR=static/js/dist/vite
+VITE_OUTDIR=static/js/dist/vite
 ENVIRONMENT=devel
 FLASK_DEBUG=true
 DEVEL=True

templates/_base-layout.html

@@ -23,10 +23,6 @@
 
     {{ vite_import('static/sass/styles.scss') }}
 
-    {% if IS_DEVELOPMENT %}
-      {{ vite_dev_tools() }}
-    {% endif %}
-
     {{ vite_import('static/js/base/base.ts') }}
     {% block scripts_includes %}{% endblock %}
     <script src="https://assets.ubuntu.com/v1/703e23c9-lazysizes+noscript+native-loading.5.1.2.min.js" defer></script>

tests/tests_vite_integration.py

@@ -4,13 +4,21 @@
 from typing import cast
 from urllib.parse import urlparse
 from pathlib import Path
+from unittest.mock import patch
+
+from webapp.app import create_app
 from webapp.vite_integration.impl import (
     ProdViteIntegration,
     DevViteIntegration,
 )
 import webapp.vite_integration.exceptions as vite_exceptions
 
-MOCK_OUTPUT_PATH = "/tmp/python_vite_test"
+
+MOCK_CONFIG = {
+    "mode": "development",
+    "port": 9999,
+    "outdir": "/tmp/python_vite_test",
+}
 MOCK_ASSET_PATH = "test/path/for/asset.ts"
 MOCK_SCSS_PATH = "test/path/for/styles.scss"
 MOCK_MANIFEST = {
@@ -41,12 +49,7 @@
 
 class TestsDevViteIntegration(TestCase):
     def setUp(self):
-        self.vite = DevViteIntegration()
-
-    def tests_dev_tools(self):
-        dev_tools = self.vite.get_dev_tools()
-        assert "@vite/client" in dev_tools
-        assert "@react-refresh" in dev_tools
+        self.vite = DevViteIntegration(MOCK_CONFIG)
 
     def tests_get_asset_url(self):
         url = self.vite.get_asset_url(MOCK_ASSET_PATH)
@@ -69,27 +72,24 @@ def tests_get_imported_css(self):
 class TestsProdViteIntegration(TestCase):
     def setUp(self):
         # create a fake Vite output directory
-        manifest_path = Path(f"{MOCK_OUTPUT_PATH}/.vite/manifest.json")
+        manifest_path = Path(f'{MOCK_CONFIG["outdir"]}/.vite/manifest.json')
         manifest_path.parent.mkdir(exist_ok=True, parents=True)
         with manifest_path.open("w+") as file:
             file.write(json.dumps(MOCK_MANIFEST))
 
         for entry in MOCK_MANIFEST.values():
             file = cast(dict, entry).get("file", "")
-            file_path = Path(f"{MOCK_OUTPUT_PATH}/{file}")
+            file_path = Path(f'{MOCK_CONFIG["outdir"]}/{file}')
             file_path.parent.mkdir(exist_ok=True, parents=True)
             with file_path.open("w+") as file:
                 file.write("")
 
-        # inject the mocks in the static class scope before initalizing
-        ProdViteIntegration.OUT_DIR = MOCK_OUTPUT_PATH
-
     def tearDown(self):
-        rmtree(MOCK_OUTPUT_PATH)
+        rmtree(MOCK_CONFIG["outdir"])
 
     def tests_good_manifest_file(self):
         # attempt to init
-        ProdViteIntegration()
+        ProdViteIntegration(MOCK_CONFIG)
 
     def tests_bad_manifest_file(self):
         # try to init a ProdViteIntegration instance with a bad manifest file
@@ -100,50 +100,44 @@ def tests_bad_manifest_file(self):
         ProdViteIntegration.BUILD_MANIFEST = "file/that/does/not/exist"
 
         with self.assertRaises(vite_exceptions.ManifestPathException):
-            self.vite = ProdViteIntegration()
+            self.vite = ProdViteIntegration(MOCK_CONFIG)
 
+        # reset build manifest path to previous value
         ProdViteIntegration.BUILD_MANIFEST = old_manifest_name
 
-    def tests_dev_tools(self):
-        vite = ProdViteIntegration()
-        dev_tools = vite.get_dev_tools()
-        assert dev_tools == ""
-
     def tests_get_asset_url__bad_asset(self):
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration(MOCK_CONFIG)
         with self.assertRaises(vite_exceptions.ManifestContentException):
             vite.get_asset_url("this_asset_does_not_exist.ts")
 
     def tests_get_asset_url__bad_path(self):
         # try to load an asset declared in the manifest but without a real
         # file backing it
+
         # load a proper manifest...
         ProdViteIntegration.manifest = MOCK_MANIFEST
         # but also load a broken OUT_DIR path
-        ProdViteIntegration.OUT_DIR = "/tmp/path/does/not/exist"
-
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration({"outdir": "/tmp/path/does/not/exist"})
         with self.assertRaises(vite_exceptions.AssetPathException):
             vite.get_asset_url(MOCK_ASSET_PATH)
 
         # cleanup
-        ProdViteIntegration.OUT_DIR = MOCK_OUTPUT_PATH
         ProdViteIntegration.manifest = None
 
     def tests_get_asset_url__is_not_ts(self):
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration(MOCK_CONFIG)
         url = vite.get_asset_url(MOCK_ASSET_PATH)
         assert MOCK_ASSET_PATH not in url  # source asset is a .ts file
         assert url.endswith(".js")  # dist asset is a .js file
 
     def tests_get_asset_url__is_not_scss(self):
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration(MOCK_CONFIG)
         url = vite.get_asset_url(MOCK_SCSS_PATH)
         assert MOCK_SCSS_PATH not in url  # source asset is a .scss file
         assert url.endswith(".css")  # dist asset is a .css file
 
     def tests_get_imported_chunks__bad_asset(self):
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration(MOCK_CONFIG)
         with self.assertRaises(vite_exceptions.ManifestContentException):
             vite.get_imported_chunks("this_asset_does_not_exist.ts")
 
@@ -153,18 +147,15 @@ def tests_get_imported_chunks__bad_path(self):
         # load a proper manifest...
         ProdViteIntegration.manifest = MOCK_MANIFEST
         # but also load a broken OUT_DIR path
-        ProdViteIntegration.OUT_DIR = "/tmp/path/does/not/exist"
-
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration({"outdir": "/tmp/path/does/not/exist"})
         with self.assertRaises(vite_exceptions.AssetPathException):
             vite.get_imported_chunks(MOCK_ASSET_PATH)
 
         # cleanup
-        ProdViteIntegration.OUT_DIR = MOCK_OUTPUT_PATH
         ProdViteIntegration.manifest = None
 
     def tests_get_imported_chunks(self):
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration(MOCK_CONFIG)
         js_entries = filter(
             lambda x: x["file"].endswith(".js"), MOCK_MANIFEST.values()
         )
@@ -173,5 +164,30 @@ def tests_get_imported_chunks(self):
         )
 
     def tests_get_imported_css(self):
-        vite = ProdViteIntegration()
+        vite = ProdViteIntegration(MOCK_CONFIG)
         assert len(vite.get_imported_css(MOCK_ASSET_PATH)) == 1
+
+
+class TestsFlaskViteExtension(TestCase):
+    @patch("webapp.config.VITE_MODE", "development")
+    @patch("webapp.config.VITE_PORT", MOCK_CONFIG["port"])
+    @patch("webapp.config.VITE_OUTDIR", MOCK_CONFIG["outdir"])
+    def setUp(self):
+        self.app = create_app(testing=True)
+        self.client = self.app.test_client()
+
+    def test_extension_init(self):
+        self.assertEqual(self.app.config.get("VITE_MODE"), "development")
+        self.assertEqual(self.app.config.get("VITE_PORT"), MOCK_CONFIG["port"])
+        self.assertEqual(
+            self.app.config.get("VITE_OUTDIR"), MOCK_CONFIG["outdir"]
+        )
+
+    def test_dev_tools(self):
+        port = self.app.config.get("VITE_PORT")
+        response = self.client.get("/")
+        self.assertEqual(response.status_code, 200)
+
+        body = response.get_data(as_text=True)
+        self.assertGreater(len(body), 0)
+        self.assertIn(f"http://localhost:{port}/@vite/client", body)

vite.config.js

@@ -107,7 +107,7 @@ export default defineConfig({
     modulePreload: false,
     emptyOutDir: true,
     sourcemap: "hidden",
-    outDir: env?.VITE_OUTPUT_DIR || "static/js/dist/vite",
+    outDir: env?.VITE_OUTDIR || "static/js/dist/vite",
     rollupOptions: {
       output: {
         entryFileNames: `[name]--[hash].js`,

webapp/app.py

@@ -13,7 +13,7 @@
 from canonicalwebteam.flask_base.app import FlaskBase
 from webapp.blog.views import init_blog
 from webapp.docs.views import init_docs
-from webapp.extensions import csrf
+from webapp.extensions import csrf, vite
 from webapp.handlers import set_handlers
 from webapp.login.views import login
 from webapp.login.oauth_views import oauth
@@ -53,15 +53,7 @@ def create_app(testing=False):
     app.name = "snapcraft"
     app.testing = testing
 
-    if not testing:
-        init_extensions(app)
-
-    if testing:
-
-        @app.context_processor
-        def inject_csrf_token():
-            return dict(csrf_token=lambda: "mocked_csrf_token")
-
+    init_extensions(app)
     set_handlers(app)
 
     app.register_blueprint(snapcraft_blueprint())
@@ -90,5 +82,13 @@ def inject_csrf_token():
     return app
 
 
-def init_extensions(app):
-    csrf.init_app(app)
+def init_extensions(app: FlaskBase):
+    vite.init_app(app)
+
+    if not app.testing:
+        csrf.init_app(app)
+    else:
+        # add a helper for injecting a mock CSRF token
+        @app.context_processor
+        def inject_csrf_token():
+            return dict(csrf_token=lambda: "mocked_csrf_token")

webapp/config.py

@@ -21,8 +21,11 @@ class ConfigurationError(Exception):
 SENTRY_DSN = os.getenv("SENTRY_DSN", "").strip()
 SENTRY_CONFIG = {"release": COMMIT_ID, "environment": ENVIRONMENT}
 DNS_VERIFICATION_SALT = os.getenv("DNS_VERIFICATION_SALT")
+
+# Vite integration config values
+VITE_MODE = "development" if IS_DEVELOPMENT else "production"
 VITE_PORT = os.getenv("VITE_PORT", 5173)
-VITE_OUTPUT_DIR = os.getenv("VITE_OUTPUT_DIR", "static/js/dist/vite")
+VITE_OUTDIR = os.getenv("VITE_OUTDIR", "static/js/dist/vite")
 
 if ENVIRONMENT != "devel":
     SESSION_COOKIE_SAMESITE = "None"

webapp/extensions.py

@@ -1,4 +1,5 @@
 from flask_wtf.csrf import CSRFProtect
-
+from webapp.vite_integration import FlaskVite
 
 csrf = CSRFProtect()
+vite = FlaskVite()

webapp/handlers.py

@@ -179,8 +179,6 @@ def snapcraft_utility_processor():
         "join": template_utils.join,
         "static_url": template_utils.static_url,
         "IS_DEVELOPMENT": IS_DEVELOPMENT,
-        "vite_import": template_utils.vite_import,
-        "vite_dev_tools": template_utils.vite_dev_tools,
         "format_number": template_utils.format_number,
         "format_display_name": template_utils.format_display_name,
         "display_name": template_utils.display_name,

webapp/template_utils.py

@@ -2,11 +2,8 @@
 import hashlib
 import os
 from dateutil import parser
-from markupsafe import Markup
 from emoji import replace_emoji
 
-from webapp.vite_integration import ViteIntegration
-
 
 # generator functions for templates
 def generate_slug(path):
@@ -90,38 +87,6 @@ def static_url(filename):
     return url + "?v=" + file_hash.hexdigest()[:7]
 
 
-def vite_import(entrypoint: str):
-    """
-    Template function that takes a .js/.ts source file as an argument and
-    returns the <script> tags with the correct src URL based on Vite's
-    output (a localhost URL in dev mode, or a static url in prod mode)
-    """
-    entry_url = ViteIntegration.get_asset_url(entrypoint)
-    is_css = entry_url.endswith((".css"))
-
-    if is_css:
-        return Markup(f'<link rel="stylesheet" href="{entry_url}" />')
-
-    entry_script = f'<script type="module" src="{entry_url}"></script>'
-
-    chunks_urls = ViteIntegration.get_imported_chunks(entrypoint)
-    chunks_scripts = [
-        f'<link rel="modulepreload" href="{c}" />' for c in chunks_urls
-    ]
-    css_urls = ViteIntegration.get_imported_css(entrypoint)
-    css_scripts = [f'<link rel="stylesheet" href="{c}" />' for c in css_urls]
-
-    return Markup(entry_script + "".join(chunks_scripts + css_scripts))
-
-
-def vite_dev_tools():
-    """
-    Template function that returns <script> tags for Vite's dev server
-    integration (or an empty string in prod mode)
-    """
-    return Markup(ViteIntegration.get_dev_tools())
-
-
 def install_snippet(
     package_name, default_track, lowest_risk_available, confinement
 ):

webapp/vite_integration/__init__.py

@@ -1,10 +1,4 @@
-from webapp.config import IS_DEVELOPMENT
-from webapp.vite_integration.impl import (
-    DevViteIntegration,
-    ProdViteIntegration,
-)
+from .extension import FlaskVite
 
 
-ViteIntegration = (
-    DevViteIntegration if IS_DEVELOPMENT else ProdViteIntegration
-)()
+__all__ = ["FlaskVite"]