namespace Moq.Context {

	public static class DbContextMockExtension
	{
		public static Mock<AHContext> SetupPartOfTheContextWith<T>(this Mock<AHContext> mockContext, List<T> listOfT, Expression<Func<AHContext, DbSet<T>>> setupReturn) where T : BaseEntity
		{
			var queryable = listOfT.AsQueryable();

			var mockSet = new Mock<DbSet<T>>();
			mockSet.Setup(x => x.Add(It.IsAny<T>())).Callback((T entity) =>
			{
				listOfT.Add(entity);
				entity.Id = listOfT.Max(x => x.Id) + 1;
			});
			mockSet.Setup(x => x.AddRange(It.IsAny<IEnumerable<T>>())).Callback((IEnumerable<T> entities) =>
			{
				var collection = entities as List<T> ?? entities.ToList();
				listOfT.AddRange(collection);
				collection.ForEach(entity => entity.Id = listOfT.Max(x => x.Id) + 1);
			});
			mockSet.As<IQueryable<T>>().Setup(m => m.Provider).Returns(queryable.Provider);
			mockSet.As<IQueryable<T>>().Setup(m => m.Expression).Returns(queryable.Expression);
			mockSet.As<IQueryable<T>>().Setup(m => m.ElementType).Returns(queryable.ElementType);
			mockSet.As<IQueryable<T>>().Setup(m => m.GetEnumerator()).Returns(queryable.GetEnumerator());

		    mockContext.Setup(x => x.Set<T>()).Returns(mockSet.Object);
			mockContext.Setup(setupReturn).Returns(mockSet.Object);

			return mockContext;
		}
	}
}