Wrapping IQueryable objects for testability

In some cases, testing IQueryable objects can be difficult. They are hard to mock and the LINQ magic that's happening under the hood is usually hidden from the programmer. It would be convenient to wrap an IQueryable around a mockable container so that our code can be tested even when the actual data sources may not be present.

To this effect, I've found that it's very easy to wrap an IQueryable around a LINQ provider that simply defers execution of the queries to the instance it wraps. The instance can then be mocked in our tests and we can continue development without this dependency. Such a wrapper provider looks much like this code.

Code:

     public class WrapperProvider : IQueryProvider
    {
        private readonly IQueryable _query;

        public WrapperProvider(IQueryable query)
        {
            _query = query;
        }

        public IQueryable<TElement> CreateQuery<TElement>(Expression expression)
        {
            return new WrapperContext<TElement>(this, expression);
        }

        public IQueryable CreateQuery(Expression expression)
        {
            return new WrapperContext(this, expression);
        }

        public TResult Execute<TResult>(Expression expression)
        {
            var newExpression = new ExpressionTreeModifier(_query).Accept(expression);

            if ((typeof(TResult).Name.StartsWith("IEnumerable")))
            {
                return (TResult)_query.Provider.CreateQuery(newExpression);
            }
            
            return (TResult)_query.Provider.Execute(newExpression);
        }

        public object Execute(Expression expression)
        {
            var newExpression = new ExpressionTreeModifier(_query).Accept(expression);
            return _query.Provider.Execute(newExpression);
        }
    }

    internal class ExpressionTreeModifier : ExpressionVisitor
    {
        private readonly IQueryable _query;

        internal ExpressionTreeModifier(IQueryable query)
        {
            _query = query;
        }

        protected override Expression VisitConstant(ConstantExpression c)
        {
            var wrapperType = typeof (WrapperContext);
            for (var iteratorType = c.Type; iteratorType != null; iteratorType = iteratorType.BaseType)
            {
                if (iteratorType == wrapperType)
                {
                    return Expression.Constant(_query);
                }
            }
            return c;
        }

        internal Expression Accept(Expression expression)
        {
            return Visit(expression);
        }
    }

    public class WrapperContext : IQueryable
    {
        public WrapperContext(IQueryable query)
        {
            Expression = Expression.Constant(this);
            Provider = new WrapperProvider(query);
        }

        internal WrapperContext(IQueryProvider provider, Expression expression)
        {
            Expression = expression;
            Provider = provider;
        }

        public IEnumerator GetEnumerator()
        {
            return Provider.Execute<IEnumerable>(Expression).GetEnumerator();
        }

        public Expression Expression { get; private set; }
        public Type ElementType { get; private set; }
        public IQueryProvider Provider { get; private set; }
    }

    public class WrapperContext<T> : WrapperContext, IQueryable<T>
    {
        public WrapperContext(IQueryable<T> query)
            : base(query)
        {
        }

        internal WrapperContext(IQueryProvider provider, Expression expression)
            : base(provider, expression)
        {
        }

        public new IEnumerator<T> GetEnumerator()
        {
            return Provider.Execute<IEnumerable<T>>(Expression).GetEnumerator();
        }
    }

In our actual code we may choose to do things like Assert.Equals(new WrapperContext<Customer>(context.Customers).Where(c => c.CustomerID == 1).Count(), 1), knowing that the execution will be deferred to whatever implementation of IQueryable the wrapper contains.

David.