From 5b9cb4a13caaccf7b9a795543143d1c8d685c0b5 Mon Sep 17 00:00:00 2001
From: Ben Broderick Phillips <bebroder@microsoft.com>
Date: Tue, 13 Feb 2024 20:32:54 -0500
Subject: [PATCH] Improve type and null handling

---
 .../job-matrix-functions.modification.tests.ps1    | 14 +++++++++++---
 .../scripts/job-matrix/job-matrix-functions.ps1    | 10 ++++++++--
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/eng/common-tests/matrix-generator/tests/job-matrix-functions.modification.tests.ps1 b/eng/common-tests/matrix-generator/tests/job-matrix-functions.modification.tests.ps1
index 9f6b5f4f719..3e1e27a453c 100644
--- a/eng/common-tests/matrix-generator/tests/job-matrix-functions.modification.tests.ps1
+++ b/eng/common-tests/matrix-generator/tests/job-matrix-functions.modification.tests.ps1
@@ -591,26 +591,34 @@ Describe "Platform Matrix Environment Variables" -Tag "UnitTest", "envvar" {
   "matrix": {
     "foo": "bar",
     "envReference": ["env:TestMatrixEnvReference", "env:TestMatrixEnvReference2", "noref"]
-  }
+  },
+  "include": [
+    {
+      "foo": "bar",
+      "envReference": "env:TestMatrixEnvReference"
+    }
+  ]
 }
 '@
 
         [System.Environment]::SetEnvironmentVariable("TestMatrixEnvReference", "")
         [array]$matrix = GenerateMatrix (GetMatrixConfigFromJson $matrixJson) "sparse"
-        $matrix.Length | Should -Be 3
+        $matrix.Length | Should -Be 4
         $matrix[0].name | Should -Be "bar"
         $matrix[0].parameters.envReference | Should -Be ""
 
         [System.Environment]::SetEnvironmentVariable("TestMatrixEnvReference", "replaced")
         [System.Environment]::SetEnvironmentVariable("TestMatrixEnvReference2", "replaced2")
         [array]$replacedMatrix = GenerateMatrix (GetMatrixConfigFromJson $matrixJson) "sparse"
-        $replacedMatrix.Length | Should -Be 3
+        $replacedMatrix.Length | Should -Be 4
         $replacedMatrix[0].name | Should -Be "bar_replaced"
         $replacedMatrix[0].parameters.envReference | Should -Be "replaced"
         $replacedMatrix[1].name | Should -Be "bar_replaced2"
         $replacedMatrix[1].parameters.envReference | Should -Be "replaced2"
         $replacedMatrix[2].name | Should -Be "bar_noref"
         $replacedMatrix[2].parameters.envReference | Should -Be "noref"
+        $replacedMatrix[3].name | Should -Be "bar_replaced"
+        $replacedMatrix[3].parameters.envReference | Should -Be "replaced"
     }
 
     It "Should support filter/replace with variable reference syntax" {
diff --git a/eng/common/scripts/job-matrix/job-matrix-functions.ps1 b/eng/common/scripts/job-matrix/job-matrix-functions.ps1
index 983baaf0a91..c11a318b9bd 100644
--- a/eng/common/scripts/job-matrix/job-matrix-functions.ps1
+++ b/eng/common/scripts/job-matrix/job-matrix-functions.ps1
@@ -76,7 +76,7 @@ class MatrixParameter {
             $displayName = $this.Value.ToString()
         }
 
-        if ($displayNamesLookup.ContainsKey($displayName)) {
+        if ($displayNamesLookup -and $displayNamesLookup.ContainsKey($displayName)) {
             $displayName = $displayNamesLookup[$displayName]
         }
 
@@ -339,6 +339,9 @@ function ProcessReplace {
 
     foreach ($element in $matrix) {
         $replacement = [MatrixParameter[]]@()
+        if (!$element -or $element.Count -eq 0) {
+            continue
+        }
 
         foreach ($perm in $element._permutation) {
             $replace = $perm
@@ -372,11 +375,14 @@ function ProcessEnvironmentVariableReferences([array]$matrix, $displayNamesLooku
 
     foreach ($element in $matrix) {
         $updated = [MatrixParameter[]]@()
+        if (!$element -or $element.Count -eq 0) {
+            continue
+        }
 
         foreach ($perm in $element._permutation) {
             # Iterate nested permutations or run once for singular values (int, string, bool)
             foreach ($flattened in $perm.Flatten()) {
-                if ($flattened.Value?.GetType() -eq "".GetType() -and $flattened.Value.StartsWith("env:")) {
+                if ($flattened.Value -is [string] -and $flattened.Value.StartsWith("env:")) {
                     $envKey = $flattened.Value.Replace("env:", "")
                     $value = [System.Environment]::GetEnvironmentVariable($envKey) ?? ""
                     if (!$value) {