diff --git a/test/Microsoft.AspNetCore.Diagnostics.EntityFrameworkCore.Tests/TestHelperExtensions.cs b/test/Microsoft.AspNetCore.Diagnostics.EntityFrameworkCore.Tests/TestHelperExtensions.cs
index d101ba3665..e87b4ef99d 100644
--- a/test/Microsoft.AspNetCore.Diagnostics.EntityFrameworkCore.Tests/TestHelperExtensions.cs
+++ b/test/Microsoft.AspNetCore.Diagnostics.EntityFrameworkCore.Tests/TestHelperExtensions.cs
@@ -8,9 +8,9 @@ namespace Microsoft.EntityFrameworkCore.Tests
 {
     public static class TestHelperExtensions
     {
-        public static EntityFrameworkServicesBuilder AddProviderServices(this EntityFrameworkServicesBuilder entityServicesBuilder)
+        public static IServiceCollection AddProviderServices(this IServiceCollection serviceCollection)
         {
-            return entityServicesBuilder.AddInMemoryDatabase();
+            return serviceCollection.AddEntityFrameworkInMemoryDatabase();
         }
 
         public static DbContextOptions UseProviderOptions(this DbContextOptions options)