Test built extension wheels in CI

scripts/ci/test_source.py

@@ -9,12 +9,15 @@
 
 import logging
 import os
+from pathlib import Path
+import subprocess
 import sys
 import tempfile
 import shutil
 import shlex
 
 from subprocess import check_output, CalledProcessError, run
+from typing import Optional
 from util import SRC_PATH
 
 logger = logging.getLogger(__name__)
@@ -64,7 +67,7 @@ def run_command(cmd, check_return_code=False, cwd=None):
         raise RuntimeError(f"{cmd} failed")
 
 
-def test_extension():
+def test_extension(whl_dir: Optional[Path] = None):
     for pkg_name, ext_path in ALL_TESTS:
         ext_name = ext_path.split('/')[-1]
         logger.info(f'installing extension: {ext_name}')
@@ -86,22 +89,50 @@ def test_extension():
         cmd = ['azdev', 'extension', 'remove', ext_name]
         run_command(cmd, check_return_code=True)
 
-
-def test_source_wheels():
+        # This catches code missing in the final wheel package, but until the full
+        # set of extension tests can be run against the wheel, it will not catch
+        # missing files required at runtime.
+        if whl_dir is not None:
+            logger.info(f'installing extension wheel: {ext_name}')
+            wheel_path = next(whl_dir.glob(f"{ext_name}*.whl"), None)
+            assert wheel_path is not None, f"Wheel for extension {ext_name} not found in {whl_dir}"
+            subprocess.run(
+                ["az", "extension", "add", "-y", "-s", wheel_path],
+                check=True,
+            )
+            subprocess.run(
+                ["az", ext_name, "--help"],
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            subprocess.run(
+                ["az", "extension", "remove", "-n", ext_name],
+                check=True,
+            )
+
+
+def test_source_wheels(whl_dir: Optional[Path] = None):
     # Test we can build all sources into wheels and that metadata from the wheel is valid
     built_whl_dir = tempfile.mkdtemp()
     source_extensions = [os.path.join(SRC_PATH, n) for n in os.listdir(SRC_PATH)
                          if os.path.isdir(os.path.join(SRC_PATH, n))]
     for s in source_extensions:
+        ext_name = s.split('/')[-1]
         if not os.path.isfile(os.path.join(s, 'setup.py')):
             continue
         try:
-            check_output(['python', 'setup.py', 'bdist_wheel', '-q', '-d', built_whl_dir], cwd=s)
+            check_output(['azdev', 'extension', 'build', ext_name, '--dist-dir', built_whl_dir])
         except CalledProcessError as err:
             raise("Unable to build extension {} : {}".format(s, err))
     # Export built wheels so CI can publish them as artifacts
-    wheels_out_dir = os.environ.get('WHEELS_OUTPUT_DIR')
-    if wheels_out_dir:
+    wheels_out_dirs = [
+        os.environ.get('WHEELS_OUTPUT_DIR'),
+        str(whl_dir) if whl_dir else None,
+    ]
+    for wheels_out_dir in wheels_out_dirs:
+        if not wheels_out_dir:
+            continue
         try:
             os.makedirs(wheels_out_dir, exist_ok=True)
             for fname in os.listdir(built_whl_dir):
@@ -115,5 +146,7 @@ def test_source_wheels():
 
 
 if __name__ == '__main__':
-    test_extension()
-    test_source_wheels()
+    with tempfile.TemporaryDirectory() as whl_dir:
+        whl_dir = Path(whl_dir)
+        test_source_wheels(whl_dir=whl_dir)
+        test_extension(whl_dir=whl_dir)