Fix more flakiness in Cosmos tests

src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseCreator.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Runtime.CompilerServices;
 using Microsoft.EntityFrameworkCore.Cosmos.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
 
@@ -20,6 +21,8 @@ public class CosmosDatabaseCreator : IDatabaseCreator
     private readonly IDatabase _database;
     private readonly ICurrentDbContext _currentContext;
     private readonly IDbContextOptions _contextOptions;
+    private readonly IExecutionStrategy _executionStrategy;
+    private readonly IDiagnosticsLogger<DbLoggerCategory.Infrastructure> _logger;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -33,14 +36,18 @@ public CosmosDatabaseCreator(
         IUpdateAdapterFactory updateAdapterFactory,
         IDatabase database,
         ICurrentDbContext currentContext,
-        IDbContextOptions contextOptions)
+        IDbContextOptions contextOptions,
+        IExecutionStrategy executionStrategy,
+        IDiagnosticsLogger<DbLoggerCategory.Infrastructure> logger)
     {
         _cosmosClient = cosmosClient;
         _designTimeModel = designTimeModel;
         _updateAdapterFactory = updateAdapterFactory;
         _database = database;
         _currentContext = currentContext;
         _contextOptions = contextOptions;
+        _executionStrategy = executionStrategy;
+        _logger = logger;
     }
 
     /// <summary>
@@ -49,26 +56,56 @@ public CosmosDatabaseCreator(
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public virtual async Task<bool> EnsureCreatedAsync(CancellationToken cancellationToken = default)
+    public virtual Task<bool> EnsureCreatedAsync(CancellationToken cancellationToken = default)
     {
-        var model = _designTimeModel.Model;
-        var created = await _cosmosClient.CreateDatabaseIfNotExistsAsync(model.GetThroughput(), cancellationToken)
-            .ConfigureAwait(false);
-
-        foreach (var container in GetContainersToCreate(model))
+        if (_currentContext.Context.ChangeTracker.Entries().Any(
+                e => e.State is EntityState.Added or EntityState.Modified or EntityState.Deleted))
         {
-            created |= await _cosmosClient.CreateContainerIfNotExistsAsync(container, cancellationToken)
-                .ConfigureAwait(false);
+            _logger.EnsureCreatedWithTrackedEntitiesWarning();
         }
 
-        if (created)
-        {
-            await InsertDataAsync(cancellationToken).ConfigureAwait(false);
-        }
+        var created = new StrongBox<bool>(false);
+        var dataInserted = new StrongBox<bool>(false);
+        var retrying = new StrongBox<bool>(false);
+        return _executionStrategy.ExecuteAsync(
+            (Creator: this, Created: created, DataInserted: dataInserted, Retrying: retrying),
+            static async (_, state, ct) =>
+            {
+                var creator = state.Creator;
+
+                if (state.Retrying.Value)
+                {
+                    creator._currentContext.Context.ChangeTracker.Clear();
+                }
+
+                state.Retrying.Value = true;
+
+                if (!state.DataInserted.Value)
+                {
+                    var model = creator._designTimeModel.Model;
+                    state.Created.Value |= await creator._cosmosClient
+                        .CreateDatabaseIfNotExistsAsync(model.GetThroughput(), ct)
+                        .ConfigureAwait(false);
+
+                    foreach (var container in GetContainersToCreate(model))
+                    {
+                        state.Created.Value |= await creator._cosmosClient
+                            .CreateContainerIfNotExistsAsync(container, ct)
+                            .ConfigureAwait(false);
+                    }
+
+                    if (state.Created.Value)
+                    {
+                        await creator.InsertDataAsync(ct).ConfigureAwait(false);
+                        state.DataInserted.Value = true;
+                    }
+                }
 
-        await SeedDataAsync(created, cancellationToken).ConfigureAwait(false);
+                await creator.SeedDataAsync(state.Created.Value, cancellationToken: ct)
+                    .ConfigureAwait(false);
 
-        return created;
+                return state.Created.Value;
+            }, verifySucceeded: null, cancellationToken);
     }
 
     private static IEnumerable<ContainerProperties> GetContainersToCreate(IModel model)
@@ -193,17 +230,17 @@ private IUpdateAdapter AddModelData()
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public virtual async Task SeedDataAsync(bool created, CancellationToken cancellationToken = default)
+    public virtual async Task SeedDataAsync(
+        bool created, CancellationToken cancellationToken = default)
     {
         var coreOptionsExtension =
-            _contextOptions.FindExtension<CoreOptionsExtension>()
-            ?? new CoreOptionsExtension();
+            _contextOptions.FindExtension<CoreOptionsExtension>();
 
-        if (coreOptionsExtension.AsyncSeeder is not null)
+        if (coreOptionsExtension?.AsyncSeeder is not null)
         {
             await coreOptionsExtension.AsyncSeeder(_currentContext.Context, created, cancellationToken).ConfigureAwait(false);
         }
-        else if (coreOptionsExtension.Seeder != null)
+        else if (coreOptionsExtension?.Seeder is not null)
         {
             throw new InvalidOperationException(CoreStrings.MissingSeeder);
         }

src/EFCore.Relational/EFCore.Relational.baseline.json

@@ -7304,56 +7304,6 @@
         }
       ]
     },
-    {
-      "Type": "class Microsoft.EntityFrameworkCore.Infrastructure.StructuredJsonPath",
-      "Methods": [
-        {
-          "Member": "StructuredJsonPath(System.Collections.Generic.IReadOnlyList<Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment> segments, int[] indices);"
-        },
-        {
-          "Member": "virtual System.Text.StringBuilder AppendTo(System.Text.StringBuilder builder);"
-        },
-        {
-          "Member": "override string ToString();"
-        }
-      ],
-      "Properties": [
-        {
-          "Member": "virtual bool IsRoot { get; }"
-        },
-        {
-          "Member": "virtual int[] Indices { get; }"
-        },
-        {
-          "Member": "static Microsoft.EntityFrameworkCore.Infrastructure.StructuredJsonPath Root { get; }"
-        },
-        {
-          "Member": "virtual System.Collections.Generic.IReadOnlyList<Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment> Segments { get; }"
-        }
-      ]
-    },
-    {
-      "Type": "sealed class Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment",
-      "Methods": [
-        {
-          "Member": "StructuredJsonPathSegment(string propertyName);"
-        },
-        {
-          "Member": "override string ToString();"
-        }
-      ],
-      "Properties": [
-        {
-          "Member": "static Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment Array { get; }"
-        },
-        {
-          "Member": "bool IsArray { get; }"
-        },
-        {
-          "Member": "string? PropertyName { get; }"
-        }
-      ]
-    },
     {
       "Type": "class Microsoft.EntityFrameworkCore.Query.JsonQueryExpression : System.Linq.Expressions.Expression, Microsoft.EntityFrameworkCore.Query.IPrintableExpression",
       "Methods": [
@@ -7970,6 +7920,9 @@
         {
           "Member": "int LastCommittedCommandIndex { get; set; }"
         },
+        {
+          "Member": "bool SeedingAttempted { get; set; }"
+        },
         {
           "Member": "Microsoft.EntityFrameworkCore.Storage.IDbContextTransaction? Transaction { get; set; }"
         }
@@ -11327,6 +11280,9 @@
         {
           "Member": "Microsoft.EntityFrameworkCore.Storage.IExecutionStrategy ExecutionStrategy { get; }"
         },
+        {
+          "Member": "Microsoft.EntityFrameworkCore.Diagnostics.IDiagnosticsLogger<Microsoft.EntityFrameworkCore.DbLoggerCategory.Infrastructure> Logger { get; init; }"
+        },
         {
           "Member": "Microsoft.EntityFrameworkCore.Migrations.IMigrationCommandExecutor MigrationCommandExecutor { get; init; }"
         },
@@ -16941,7 +16897,8 @@
       ],
       "Properties": [
         {
-          "Member": "static string BadSequenceString { get; }"
+          "Member": "static string BadSequenceString { get; }",
+          "Stage": "Obsolete"
         },
         {
           "Member": "static string BadSequenceType { get; }"
@@ -19928,6 +19885,56 @@
         }
       ]
     },
+    {
+      "Type": "class Microsoft.EntityFrameworkCore.Infrastructure.StructuredJsonPath",
+      "Methods": [
+        {
+          "Member": "StructuredJsonPath(System.Collections.Generic.IReadOnlyList<Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment> segments, int[] indices);"
+        },
+        {
+          "Member": "virtual System.Text.StringBuilder AppendTo(System.Text.StringBuilder builder);"
+        },
+        {
+          "Member": "override string ToString();"
+        }
+      ],
+      "Properties": [
+        {
+          "Member": "virtual int[] Indices { get; }"
+        },
+        {
+          "Member": "virtual bool IsRoot { get; }"
+        },
+        {
+          "Member": "static Microsoft.EntityFrameworkCore.Infrastructure.StructuredJsonPath Root { get; }"
+        },
+        {
+          "Member": "virtual System.Collections.Generic.IReadOnlyList<Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment> Segments { get; }"
+        }
+      ]
+    },
+    {
+      "Type": "sealed class Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment",
+      "Methods": [
+        {
+          "Member": "StructuredJsonPathSegment(string propertyName);"
+        },
+        {
+          "Member": "override string ToString();"
+        }
+      ],
+      "Properties": [
+        {
+          "Member": "static Microsoft.EntityFrameworkCore.Metadata.StructuredJsonPathSegment Array { get; }"
+        },
+        {
+          "Member": "bool IsArray { get; }"
+        },
+        {
+          "Member": "string? PropertyName { get; }"
+        }
+      ]
+    },
     {
       "Type": "class Microsoft.EntityFrameworkCore.Metadata.Builders.TableBuilder : Microsoft.EntityFrameworkCore.Infrastructure.IInfrastructure<Microsoft.EntityFrameworkCore.Metadata.Builders.EntityTypeBuilder>",
       "Methods": [

src/EFCore.Relational/Migrations/Internal/Migrator.cs

@@ -192,6 +192,12 @@ private bool MigrateImplementation(
             var seed = coreOptionsExtension.Seeder;
             if (seed != null)
             {
+                if (state.SeedingAttempted)
+                {
+                    context.ChangeTracker.Clear();
+                }
+
+                state.SeedingAttempted = true;
                 seed(context, state.AnyOperationPerformed);
             }
             else if (coreOptionsExtension.AsyncSeeder != null)
@@ -328,6 +334,12 @@ await _migrationCommandExecutor.ExecuteNonQueryAsync(
             var seedAsync = coreOptionsExtension.AsyncSeeder;
             if (seedAsync != null)
             {
+                if (state.SeedingAttempted)
+                {
+                    context.ChangeTracker.Clear();
+                }
+
+                state.SeedingAttempted = true;
                 await seedAsync(context, state.AnyOperationPerformed, cancellationToken).ConfigureAwait(false);
             }
             else if (coreOptionsExtension.Seeder != null)

src/EFCore.Relational/Migrations/MigrationExecutionState.cs

@@ -32,4 +32,9 @@ public sealed class MigrationExecutionState
     ///     The transaction that is in use.
     /// </summary>
     public IDbContextTransaction? Transaction { get; set; }
+
+    /// <summary>
+    ///     Indicates whether seeding has been attempted.
+    /// </summary>
+    public bool SeedingAttempted { get; set; }
 }