From 0e550592b5ce207b2bb5151c792d14ffda5577b6 Mon Sep 17 00:00:00 2001
From: Kaffeetasse <7276999+Kaffeetasse@users.noreply.github.com>
Date: Sat, 14 Oct 2023 01:51:35 +0200
Subject: [PATCH] fix initial entities with predefined ids

---
 .../KeyContext.cs                             |  6 +++
 .../KeyFactoryBuilderBase.cs                  | 37 ++++++++++++++++++-
 .../DbContextMockTests.cs                     | 19 ++++++++++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/src/EntityFrameworkCoreMock.Shared/KeyContext.cs b/src/EntityFrameworkCoreMock.Shared/KeyContext.cs
index 3ba515b..ea75cfe 100644
--- a/src/EntityFrameworkCoreMock.Shared/KeyContext.cs
+++ b/src/EntityFrameworkCoreMock.Shared/KeyContext.cs
@@ -11,5 +11,11 @@ public sealed class KeyContext
         private long _nextIdentity = 1;
 
         public long NextIdentity => _nextIdentity++;
+
+        public long CurrentIdentity
+        {
+            get => _nextIdentity;
+            set => _nextIdentity = value + 1;
+        }
     }
 }
diff --git a/src/EntityFrameworkCoreMock.Shared/KeyFactoryBuilders/KeyFactoryBuilderBase.cs b/src/EntityFrameworkCoreMock.Shared/KeyFactoryBuilders/KeyFactoryBuilderBase.cs
index 14444a1..58fa019 100644
--- a/src/EntityFrameworkCoreMock.Shared/KeyFactoryBuilders/KeyFactoryBuilderBase.cs
+++ b/src/EntityFrameworkCoreMock.Shared/KeyFactoryBuilders/KeyFactoryBuilderBase.cs
@@ -37,11 +37,11 @@ private static Func<T, KeyContext, object> BuildIdentityKeyFactory<T>(PropertyIn
 
             if (keyProperty.PropertyType == typeof(int))
             {
-                return BuildIdentityKeyFactory<T, int>(keyProperty, ctx => Expression.Property(ctx, nameof(KeyContext.NextIdentity)));
+                return BuildAutoincrementIdentityKeyFactory<T, int>(keyProperty, ctx => Expression.Property(ctx, nameof(KeyContext.NextIdentity)));
             }
             else if (keyProperty.PropertyType == typeof(long))
             {
-                return BuildIdentityKeyFactory<T, long>(keyProperty, ctx => Expression.Property(ctx, nameof(KeyContext.NextIdentity)));
+                return BuildAutoincrementIdentityKeyFactory<T, long>(keyProperty, ctx => Expression.Property(ctx, nameof(KeyContext.NextIdentity)));
             }
             else if (keyProperty.PropertyType == typeof(Guid))
             {
@@ -71,6 +71,39 @@ private static Func<TEntity, KeyContext, object> BuildIdentityKeyFactory<TEntity
 
             return Expression.Lambda<Func<TEntity, KeyContext, object>>(body, entityArgument, keyContextArgument).Compile();
         }
+        
+        private static Func<TEntity, KeyContext, object> BuildAutoincrementIdentityKeyFactory<TEntity, TKey>(
+            PropertyInfo keyProperty,
+            Func<ParameterExpression, Expression> nextIdentity)
+        {
+            var entityArgument = Expression.Parameter(typeof(TEntity));
+            var keyContextArgument = Expression.Parameter(typeof(KeyContext));
+            var keyValueVariable = Expression.Variable(typeof(TKey));
+            var body = Expression.Block(typeof(object),
+                new[] { keyValueVariable },
+                Expression.Assign(keyValueVariable, Expression.Convert(Expression.Property(entityArgument, keyProperty), typeof(TKey))),
+                Expression.IfThenElse(Expression.Equal(keyValueVariable, Expression.Default(typeof(TKey))),
+                    Expression.Block(
+                        Expression.Assign(keyValueVariable, Expression.Convert(nextIdentity(keyContextArgument), typeof(TKey))),
+                        Expression.Assign(Expression.Property(entityArgument, keyProperty), keyValueVariable)
+                    ),
+                    Expression.Block(
+                        Expression.Assign(
+                            Expression.Property(keyContextArgument, nameof(KeyContext.CurrentIdentity)),
+                            Expression.Call(
+                                typeof(Math),
+                                nameof(Math.Max),
+                                Type.EmptyTypes,
+                                Expression.Convert(keyValueVariable, typeof(long)),
+                                Expression.Property(keyContextArgument, nameof(KeyContext.CurrentIdentity))
+                            )
+                        )
+                    )
+                ),
+                Expression.Convert(keyValueVariable, typeof(object)));
+
+            return Expression.Lambda<Func<TEntity, KeyContext, object>>(body, entityArgument, keyContextArgument).Compile();
+        }
 
         private static Func<T, KeyContext, object> BuildDefaultKeyFactory<T>(PropertyInfo[] keyProperties)
         {
diff --git a/tests/EntityFrameworkCoreMock.Moq.Tests/DbContextMockTests.cs b/tests/EntityFrameworkCoreMock.Moq.Tests/DbContextMockTests.cs
index 5d5846c..97bb92e 100644
--- a/tests/EntityFrameworkCoreMock.Moq.Tests/DbContextMockTests.cs
+++ b/tests/EntityFrameworkCoreMock.Moq.Tests/DbContextMockTests.cs
@@ -165,6 +165,25 @@ public void DbContextMock_CreateDbSetMock_AddMultipleModelsWithDatabaseGenerated
             Assert.That(dbSetMock.Object.First(x => x.Id == 2).Value, Is.EqualTo("second"));
             Assert.That(dbSetMock.Object.First(x => x.Id == 3).Value, Is.EqualTo("third"));
         }
+        
+        [Test]
+        public void DbContextMock_CreateDbSetMock_AddWithDatabaseGeneratedIdentityKeyWithIdsOnInitialEntities_ShouldGenerateSequentialKey()
+        {
+            var dbContextMock = new DbContextMock<TestDbContext>(Options);
+            var dbSetMock = dbContextMock.CreateDbSetMock(x => x.GeneratedKeyModels, new[]
+            {
+                new GeneratedKeyModel {Id = 1, Value = "first"},
+                new GeneratedKeyModel {Id = 2, Value = "second"}
+            });
+            dbSetMock.Object.Add(new GeneratedKeyModel { Value = "third" });
+            dbContextMock.Object.SaveChanges();
+
+            Assert.That(dbSetMock.Object.Min(x => x.Id), Is.EqualTo(1));
+            Assert.That(dbSetMock.Object.Max(x => x.Id), Is.EqualTo(3));
+            Assert.That(dbSetMock.Object.First(x => x.Id == 1).Value, Is.EqualTo("first"));
+            Assert.That(dbSetMock.Object.First(x => x.Id == 2).Value, Is.EqualTo("second"));
+            Assert.That(dbSetMock.Object.First(x => x.Id == 3).Value, Is.EqualTo("third"));
+        }
 
         [Test]
         public void DbContextMock_CreateDbSetMock_AddMultipleModelsWithGuidAsDatabaseGeneratedIdentityKey_ShouldGenerateRandomGuidAsKey()