Read env vars from _index.md

Author: federicotdn

.github/workflows/run-code-blocks.yml

@@ -35,5 +35,5 @@ jobs:
           ALL_CHANGED_FILES: ${{ steps.changed-files.outputs.all_changed_files }}
         run: |
           for file in ${ALL_CHANGED_FILES}; do
-            python -u scripts/md-k6.py --duration 2s "$file"
+            python -u scripts/md-k6.py --duration 2s "$file" --index-file docs/sources/k6/_index.md
           done

docs/sources/k6/next/javascript-api/jslib/http-instrumentation-pyroscope/_index.md

@@ -34,7 +34,7 @@ This module, by default, adds three key-value pairs:
 
 ## Example
 
-This example demonstrates how to use the this library to instrument every HTTP request made in a script with the `baggage` header.
+This example demonstrates how to use the this library to instrument every HTTP request made in a script with the `baggage` header TEST.
 
 {{< code >}}
 

scripts/md-k6.py

@@ -11,6 +11,7 @@
 import argparse
 import subprocess
 import tempfile
+from typing import TextIO
 from unittest import TestCase
 from collections import namedtuple
 
@@ -120,7 +121,11 @@ def run_k6(script: Script, duration: str | None, verbose: bool) -> None:
             print(line)
 
 
-def extract_scripts(text: str, lang: str = JS) -> list[Script]:
+def extract_scripts(
+    text: str, lang: str = JS, index_env: dict[str, str] | None = None
+) -> list[Script]:
+    index_env = index_env or {}
+
     # A somewhat complicated regex in order to make parsing of the code block
     # easier. Essentially, takes this:
     #
@@ -185,6 +190,13 @@ def extract_scripts(text: str, lang: str = JS) -> list[Script]:
             args.append(opt.removeprefix(ARG))
 
         script_text = "\n".join(lines[1:])
+        for key, value in index_env.items():
+            # Replace instances of '{{< param "FOO_BAR" >}}' with '1234', assuming
+            # the env value of FOO_BAR=1234.
+            script_text = re.sub(
+                r"{{ *< *param *\"" + key + '" *> *}}', value, script_text
+            )
+
         scripts.append(
             Script(
                 text=script_text,
@@ -198,6 +210,26 @@ def extract_scripts(text: str, lang: str = JS) -> list[Script]:
     return scripts
 
 
+def read_front_matter_env(f: TextIO) -> dict[str, str]:
+    # Read a Front Matter .md file with environment variables.
+    # The variables are actually in a YAML dictionary, but since
+    # I want to avoid 3rd party library, I will just try to parse
+    # them line-by-line.
+    result = {}
+    for line in f.readlines():
+        if ":" not in line:
+            continue
+
+        parts = line.strip().split(":")
+        key, value = parts[0].strip(), parts[1].strip()
+
+        # Look for YAML keys with ALL_CAPS_NAMES
+        if re.match(r"[A-Z0-9_]", key) and value:
+            result[key] = value
+
+    return result
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         description="Run k6 scripts within Markdown files."
@@ -215,6 +247,12 @@ def main() -> None:
         default=None,
         help="Override script(s) duration. Is not applied to browser tests.",
     )
+    parser.add_argument(
+        "--index-file",
+        type=argparse.FileType(),
+        default=None,
+        help="Path to Front Matter _index.md file containing environment variables.",
+    )
     parser.add_argument(
         "--verbose",
         "-v",
@@ -223,8 +261,13 @@ def main() -> None:
         action="store_true",
     )
     args = parser.parse_args()
+    index_env = read_front_matter_env(args.index_file)
 
     print("Starting md-k6 script.")
+    print("Front Matter variables:")
+    for key, value in index_env.items():
+        print(f"  {key}={value}")
+
     print("Reading from file:", args.file.name)
 
     lang = args.lang
@@ -234,7 +277,7 @@ def main() -> None:
         print(f"Skipping entire file ({SKIP_ALL}).")
         return
 
-    scripts = extract_scripts(text, lang)
+    scripts = extract_scripts(text, lang, index_env)
 
     if ":" in args.blocks:
         range_parts = args.blocks.split(":")
@@ -290,6 +333,17 @@ def testExtractScriptSimple(self):
         self.assertEqual(len(scripts), 1)
         self.assertEqual(scripts[0].text, "hello, world!")
 
+    def testExtractScriptIndexEnv(self):
+        text = """
+```javascript
+{{< param "FOO_BAR" >}} {{  < param    "FOO_BAR"   > }}
+```
+        """
+        scripts = extract_scripts(text, JS, {"FOO_BAR": "testing"})
+
+        self.assertEqual(len(scripts), 1)
+        self.assertEqual(scripts[0].text, "testing testing")
+
     def testExtractScriptOptions(self):
         text = """
 <!-- md-k6:nofail,env.A=B,env.C=X- -->