From e914e923d88cf07f5282179e6fedff5787b29a22 Mon Sep 17 00:00:00 2001
From: Andriy Svyryd <AndriySvyryd@users.noreply.github.com>
Date: Thu, 15 Aug 2024 10:10:23 -0700
Subject: [PATCH] [release/9.0-rc1] Fix OptimisticConcurrencyCosmosTest
 (#34443)

---
 test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs |  9 ++++-----
 test/EFCore.Specification.Tests/F1FixtureBase.cs      | 10 ++++++++--
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs b/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs
index bb18ffb7ed1..78894274ba2 100644
--- a/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs
+++ b/test/EFCore.Cosmos.FunctionalTests/F1CosmosFixture.cs
@@ -16,20 +16,19 @@ protected override ITestStoreFactory TestStoreFactory
     public override TestHelpers TestHelpers
         => CosmosTestHelpers.Instance;
 
-    public override async Task ReseedAsync()
+    protected override async Task<bool> ShouldSeedAsync(F1Context context)
     {
-        await base.ReseedAsync();
-
-        using var context = CreateContext();
         try
         {
-            await context.Teams.SingleAsync(t => t.Id == Team.Ferrari);
+            await base.ShouldSeedAsync(context);
         }
         catch
         {
             // Recreating the containers without using CosmosClient causes cached metadata in CosmosClient to be out of sync
             // and causes the first query to fail. This is a workaround for that.
         }
+
+        return await base.ShouldSeedAsync(context);
     }
 
     public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder)
diff --git a/test/EFCore.Specification.Tests/F1FixtureBase.cs b/test/EFCore.Specification.Tests/F1FixtureBase.cs
index 8da3c33f3e4..b785fcfc11e 100644
--- a/test/EFCore.Specification.Tests/F1FixtureBase.cs
+++ b/test/EFCore.Specification.Tests/F1FixtureBase.cs
@@ -19,7 +19,7 @@ public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder build
             .UseModel(CreateModelExternal())
             .UseSeeding((c, _) =>
             {
-                if (c.Set<EngineSupplier>().Count() != 0)
+                if (!ShouldSeed((F1Context)c))
                 {
                     return;
                 }
@@ -29,7 +29,7 @@ public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder build
             })
             .UseAsyncSeeding(async (c, _, t) =>
             {
-                if (await c.Set<EngineSupplier>().CountAsync(t) != 0)
+                if (!await ShouldSeedAsync((F1Context)c))
                 {
                     return;
                 }
@@ -40,6 +40,12 @@ public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder build
             .ConfigureWarnings(
                 w => w.Ignore(CoreEventId.SaveChangesStarting, CoreEventId.SaveChangesCompleted));
 
+    protected virtual bool ShouldSeed(F1Context context)
+        => context.EngineSuppliers.Count() == 0;
+
+    protected virtual async Task<bool> ShouldSeedAsync(F1Context context)
+        => await context.EngineSuppliers.CountAsync() == 0;
+
     protected override IServiceCollection AddServices(IServiceCollection serviceCollection)
         => base.AddServices(serviceCollection.AddSingleton<ISingletonInterceptor, F1MaterializationInterceptor>());