Ensure registry module restoration works as expected for import and…

… `test` statements
Azure · Jul 27, 2023 · 9a6329d · 9a6329d
1 parent 2a691da
commit 9a6329d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 10 deletions.
diff --git a/src/Bicep.Cli/Services/CompilationService.cs b/src/Bicep.Cli/Services/CompilationService.cs
@@ -7,6 +7,7 @@
 using Bicep.Core.Diagnostics;
 using Bicep.Core.Extensions;
 using Bicep.Core.FileSystem;
+using Bicep.Core.Navigation;
 using Bicep.Core.Registry;
 using Bicep.Core.Semantics;
 using Bicep.Core.Syntax;
@@ -126,14 +127,14 @@ public DecompileResult DecompileParams(string inputPath, string outputPath, stri
             {
                 workspace.UpsertSourceFile(SourceFileFactory.CreateBicepFile(fileUri, bicepOutput));
             }
-           
+
             return decompilation;
         }
 
         private static ImmutableDictionary<BicepSourceFile, ImmutableArray<IDiagnostic>> GetModuleRestoreDiagnosticsByBicepFile(SourceFileGrouping sourceFileGrouping, ImmutableHashSet<ModuleSourceResolutionInfo> originalModulesToRestore, bool forceModulesRestore)
         {
-            static IDiagnostic? DiagnosticForModule(SourceFileGrouping grouping, ModuleDeclarationSyntax module)
-                => grouping.TryGetErrorDiagnostic(module) is { } errorBuilder ? errorBuilder(DiagnosticBuilder.ForPosition(module.Path)) : null;
+            static IDiagnostic? DiagnosticForModule(SourceFileGrouping grouping, IForeignTemplateReference module)
+                => grouping.TryGetErrorDiagnostic(module) is { } errorBuilder ? errorBuilder(DiagnosticBuilder.ForPosition(module.ReferenceSourceSyntax)) : null;
 
             static IEnumerable<(BicepFile, IDiagnostic)> GetDiagnosticsForModulesToRestore(SourceFileGrouping grouping, ImmutableHashSet<ModuleSourceResolutionInfo> originalModulesToRestore)
             {

diff --git a/src/Bicep.Core/Registry/ModuleDispatcherExtensions.cs b/src/Bicep.Core/Registry/ModuleDispatcherExtensions.cs
@@ -3,7 +3,6 @@
 
 using Bicep.Core.Extensions;
 using Bicep.Core.Modules;
-using Bicep.Core.Syntax;
 using Bicep.Core.Workspaces;
 using System.Collections.Generic;
 using System.Linq;
@@ -14,7 +13,7 @@ public static class ModuleDispatcherExtensions
     {
         public static IEnumerable<ModuleReference> GetValidModuleReferences(this IModuleDispatcher moduleDispatcher, IEnumerable<ModuleSourceResolutionInfo> modules) =>
             modules
-                .Select(t => moduleDispatcher.TryGetModuleReference(t.ModuleDeclaration, t.ParentTemplateFile.FileUri, out var moduleRef, out _) ? moduleRef : null)
+                .Select(t => moduleDispatcher.TryGetModuleReference(t.ForeignTemplateReference, t.ParentTemplateFile.FileUri, out var moduleRef, out _) ? moduleRef : null)
                 .WhereNotNull();
     }
 }
diff --git a/src/Bicep.Core/Workspaces/SourceFileGrouping.cs b/src/Bicep.Core/Workspaces/SourceFileGrouping.cs
@@ -7,7 +7,6 @@
 using Bicep.Core.Extensions;
 using Bicep.Core.FileSystem;
 using Bicep.Core.Navigation;
-using Bicep.Core.Syntax;
 using static Bicep.Core.Diagnostics.DiagnosticBuilder;
 
 namespace Bicep.Core.Workspaces
@@ -23,7 +22,7 @@ public record UriResolutionResult(
         ErrorBuilderDelegate? ErrorBuilder);
 
     public record ModuleSourceResolutionInfo(
-        ModuleDeclarationSyntax ModuleDeclaration,
+        IForeignTemplateReference ForeignTemplateReference,
         ISourceFile ParentTemplateFile);
 
     public record SourceFileGrouping(
@@ -35,9 +34,9 @@ public record SourceFileGrouping(
     {
         public IEnumerable<ModuleSourceResolutionInfo> GetModulesToRestore()
             => UriResultByModule.SelectMany(
-                kvp => kvp.Value.Keys.OfType<ModuleDeclarationSyntax>()
-                    .Where(x => kvp.Value[x].RequiresRestore)
-                    .Select(mds => new ModuleSourceResolutionInfo(mds, kvp.Key)));
+                kvp => kvp.Value
+                    .Where(entry => entry.Value.RequiresRestore)
+                    .Select(entry => new ModuleSourceResolutionInfo(entry.Key, kvp.Key)));
 
         public BicepSourceFile EntryPoint => (FileResultByUri[EntryFileUri].File as BicepSourceFile)!;