Fix Async change tracking not working

src/MongoFramework/Infrastructure/Linq/MongoFrameworkQueryProvider.cs

@@ -234,7 +234,10 @@ private async IAsyncEnumerable<TResult> ExecuteModelAsync<TResult>(AggregateExec
 					var resultBatch = underlyingCursor.Current;
 					foreach (var item in resultBatch)
 					{
-						if (item is TEntity entityItem && (model.ResultTransformer == null || model.ResultTransformer.ReturnType == typeof(TEntity)))
+						if (item is TEntity entityItem &&
+						    (model.ResultTransformer == null ||
+							model.ResultTransformer.ReturnType == typeof(ValueTask<TEntity>) ||
+							model.ResultTransformer.ReturnType == typeof(Task<TEntity>)))
 						{
 							EntityProcessors.ProcessEntity(entityItem, Connection);
 						}

tests/MongoFramework.Tests/MongoDbSetTests.cs

@@ -3,6 +3,7 @@
 using System.ComponentModel.DataAnnotations;
 using System.Linq;
 using System.Threading.Tasks;
+using MongoFramework.Linq;
 
 namespace MongoFramework.Tests
 {
@@ -362,6 +363,60 @@ public void SuccessfullyRemoveRangeByPredicate()
 			context.SaveChanges();
 			Assert.AreEqual(1, dbSet.Count(m => m.Description == "SuccessfullyRemoveRangeByPredicate"));
 			Assert.IsNotNull(dbSet.FirstOrDefault(m => m.Id == entities[0].Id));
+		}
+
+		[TestMethod]
+		public void SuccessfullyLinqFindTracked()
+		{
+			var connection = TestConfiguration.GetConnection();
+			var context = new MongoDbContext(connection);
+			var dbSet = new MongoDbSet<TestModel>(context);
+
+			var model = new TestModel
+			{
+				Id = "abcd",
+				Description = "SuccessfullyFindTracked.1"
+			};
+
+			dbSet.Add(model);
+
+			context.SaveChanges();
+
+			ResetMongoDb();
+
+			var result = dbSet.FirstOrDefault();
+			result.Description = "changed";
+			context.ChangeTracker.DetectChanges();
+
+			Assert.AreEqual(MongoFramework.Infrastructure.EntityEntryState.Updated, context.ChangeTracker.GetEntry(result).State);
 		}
-	}
+
+		[TestMethod]
+		public async Task SuccessfullyLinqFindTrackedAsync()
+		{
+			var connection = TestConfiguration.GetConnection();			
+			var context = new MongoDbContext(connection);
+			var dbSet = new MongoDbSet<TestModel>(context);
+
+			var model = new TestModel
+			{
+				Id = "abcd",
+				Description = "SuccessfullyFindTracked.1"
+			};
+
+			dbSet.Add(model);
+
+			context.SaveChanges();
+
+			ResetMongoDb();
+
+			var result = await dbSet.FirstOrDefaultAsync();
+			result.Description = "changed";
+			context.ChangeTracker.DetectChanges();
+
+			Assert.AreEqual(MongoFramework.Infrastructure.EntityEntryState.Updated, context.ChangeTracker.GetEntry(result).State);
+		}
+
+	}
+
 }

tests/MongoFramework.Tests/MongoDbTenantSetTests.cs

@@ -5,7 +5,9 @@
 using System.Threading.Tasks;
 using MongoDB.Driver;
 using MongoFramework.Attributes;
-
+using System.Runtime.CompilerServices;
+using MongoFramework.Linq;
+
 namespace MongoFramework.Tests
 {
 	[TestClass]
@@ -819,6 +821,59 @@ public void BlocksDuplicatesByTenant()
 			dbSet.Add(new TestUniqueModel{UserName = "BlocksDuplicatesByTenant"});
 			Assert.ThrowsException<MongoBulkWriteException<TestUniqueModel>>(() => context.SaveChanges());
 		}
+		[TestMethod]
+		public void SuccessfullyLinqFindTracked()
+		{
+			var connection = TestConfiguration.GetConnection();
+			var tenantId = TestConfiguration.GetTenantId();
+			var context = new MongoDbTenantContext(connection, tenantId);
+			var dbSet = new MongoDbTenantSet<TestModel>(context);
+
+			var model = new TestModel
+			{
+				Id = "abcd",
+				Description = "SuccessfullyFindTracked.1"
+			};
+
+			dbSet.Add(model);
+
+			context.SaveChanges();
+
+			ResetMongoDb();
+
+			var result = dbSet.FirstOrDefault();
+			result.Description = "changed";
+			context.ChangeTracker.DetectChanges();
+
+			Assert.AreEqual(MongoFramework.Infrastructure.EntityEntryState.Updated, context.ChangeTracker.GetEntry(result).State);
+		}
+
+		[TestMethod]
+		public async Task SuccessfullyLinqFindTrackedAsync()
+		{
+			var connection = TestConfiguration.GetConnection();
+			var tenantId = TestConfiguration.GetTenantId();
+			var context = new MongoDbTenantContext(connection, tenantId);
+			var dbSet = new MongoDbTenantSet<TestModel>(context);
+
+			var model = new TestModel
+			{
+				Id = "abcd",
+				Description = "SuccessfullyFindTracked.1"
+			};
+
+			dbSet.Add(model);
+
+			context.SaveChanges();
+
+			ResetMongoDb();
+
+			var result = await dbSet.FirstOrDefaultAsync();
+			result.Description = "changed";
+			context.ChangeTracker.DetectChanges();
+
+			Assert.AreEqual(MongoFramework.Infrastructure.EntityEntryState.Updated, context.ChangeTracker.GetEntry(result).State);
+		}
 
 	}
 }