Ensure registry module restoration works as expected for import and…

… `test` statements

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)
  {

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();
  }
 }

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)!;