From 8cdcef0d0037ff32dc22c721d4b3835ccabc582f Mon Sep 17 00:00:00 2001
From: Mikhail Timofeev <48208649+miketimofeev@users.noreply.github.com>
Date: Mon, 8 Mar 2021 22:43:58 +0300
Subject: [PATCH] Initialize ChocolateyToolsLocation variable and fix tests
 (#2874)

---
 images/win/scripts/Installers/Initialize-VM.ps1   | 4 ++++
 images/win/scripts/Installers/Install-Haskell.ps1 | 6 ++++++
 images/win/scripts/Tests/Haskell.Tests.ps1        | 8 +++++++-
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/images/win/scripts/Installers/Initialize-VM.ps1 b/images/win/scripts/Installers/Initialize-VM.ps1
index 7623497e..ff02d065 100644
--- a/images/win/scripts/Installers/Initialize-VM.ps1
+++ b/images/win/scripts/Installers/Initialize-VM.ps1
@@ -132,6 +132,10 @@ if (Test-IsWin19) {
     Choco-Install -PackageName vcredist2010
 }
 
+# Initialize environmental variable ChocolateyToolsLocation by invoking choco Get-ToolsLocation function
+Import-Module "$env:ChocolateyInstall\helpers\chocolateyInstaller.psm1" -Force
+Get-ToolsLocation
+
 # Expand disk size of OS drive
 $driveLetter = "C"
 $size = Get-PartitionSupportedSize -DriveLetter $driveLetter
diff --git a/images/win/scripts/Installers/Install-Haskell.ps1 b/images/win/scripts/Installers/Install-Haskell.ps1
index 18b256a3..e34d930b 100644
--- a/images/win/scripts/Installers/Install-Haskell.ps1
+++ b/images/win/scripts/Installers/Install-Haskell.ps1
@@ -19,6 +19,12 @@ ForEach ($version in $VersionsList)
 $DefaultGhcVersion = $VersionsList | Select-Object -Last 1
 $DefaultGhcShortVersion = ([version]$DefaultGhcVersion).ToString(3)
 $DefaultGhcPath = Join-Path $env:ChocolateyInstall "lib\ghc.$DefaultGhcVersion\tools\ghc-$DefaultGhcShortVersion\bin"
+# Starting from version 9 haskell installation directory is $env:ChocolateyToolsLocation instead of $env:ChocolateyInstall\lib
+if ($ghcVersion -notmatch "^[0-8]\.\d+.*")
+{
+    $DefaultGhcPath = Join-Path $env:ChocolateyToolsLocation "ghc-$DefaultGhcShortVersion\bin"
+}
+
 Add-MachinePathItem -PathItem $DefaultGhcPath
 
 Write-Host "Installing cabal..."
diff --git a/images/win/scripts/Tests/Haskell.Tests.ps1 b/images/win/scripts/Tests/Haskell.Tests.ps1
index 826dc16d..4bddc57a 100644
--- a/images/win/scripts/Tests/Haskell.Tests.ps1
+++ b/images/win/scripts/Tests/Haskell.Tests.ps1
@@ -11,10 +11,16 @@ Describe "Haskell" {
     $ghcTestCases = $ghcVersionList | ForEach-Object {
         $ghcVersion = $_
         $ghcShortVersion = ([version]$ghcVersion).ToString(3)
+        $binGhcPath = Join-Path $chocoPackagesPath "ghc.$ghcVersion\tools\ghc-$ghcShortVersion\bin\ghc.exe"
+        # Starting from version 9 haskell installation directory is $env:ChocolateyToolsLocation instead of $env:ChocolateyInstall\lib
+        if ($ghcVersion -notmatch "^[0-8]\.\d+.*")
+        {
+            $binGhcPath = Join-Path $env:ChocolateyToolsLocation "ghc-$ghcShortVersion\bin\ghc.exe"
+        }
         @{
             ghcVersion = $ghcVersion
             ghcShortVersion = $ghcShortVersion
-            binGhcPath = Join-Path $chocoPackagesPath "ghc.$ghcVersion\tools\ghc-$ghcShortVersion\bin\ghc.exe"
+            binGhcPath = $binGhcPath
         }
     }