diff --git a/test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test.csproj b/test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test.csproj
index fb425c2263..95834c816a 100644
--- a/test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test.csproj
+++ b/test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test/Microsoft.AspNetCore.Mvc.Razor.Extensions.Test.csproj
@@ -30,21 +30,23 @@
     <PackageReference Include="xunit.runner.visualstudio" Version="$(XunitVersion)" />
 
     <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="$(RoslynVersion)" />
+  </ItemGroup>
 
-    <!-- We need this to compile but don't want it referenced by the project.
-      It is included dynamically in CodeGeneration tests. -->
+  <ItemGroup>
+    <!-- We need this to compile but don't want it referenced by the project. It is included dynamically in CodeGeneration tests. -->
     <ProjectReference Include="..\Microsoft.AspNetCore.Razor.Test.MvcShim\Microsoft.AspNetCore.Razor.Test.MvcShim.csproj">
       <ReferenceOutputAssembly>false</ReferenceOutputAssembly>
       <OutputItemType>Content</OutputItemType>
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
       <Targets>Build</Targets>
     </ProjectReference>
+
+    <!-- Due to how DependencyContext.Load handles our special MvcShim reference above we need to also reference the MvcShim's DiagnosticSource dependency -->
+    <PackageReference Include="System.Diagnostics.DiagnosticSource" Version="$(DiagnosticSourceVersion)" />
   </ItemGroup>
 
   <ItemGroup>
-    <MvcShim
-      Condition="$(TargetFramework) != ''"
-      Include="..\Microsoft.AspNetCore.Razor.Test.MvcShim\bin\$(Configuration)\$(TargetFramework)\Microsoft.AspNetCore.Razor.Test.MvcShim.deps.json" />
+    <MvcShim Condition="$(TargetFramework) != ''" Include="..\Microsoft.AspNetCore.Razor.Test.MvcShim\bin\$(Configuration)\$(TargetFramework)\Microsoft.AspNetCore.Razor.Test.MvcShim.deps.json" />
   </ItemGroup>
 
   <Target Name="CopyDepsFromShims" AfterTargets="Build">