Mock `TableClient.QueryAsync()` for Unit Test

I have the following method that I need to test:

public async Task<SomeClass> GetAsync(string partitionKey, string rowKey)
{
    var entities = new List<SomeClass>();
    await foreach (var e in _tableClient.QueryAsync<SomeClass>(x => x.PartitionKey == partitionKey && x.RowKey == rowKey))
    {
        entities.Add(e);
    }
    return entities.FirstOrDefault();
}

I'd like to setup the _tableClient.QueryAsync() (in the moq) to be able to return different result based on the input parameter. This is important to ensure my unit test covers the logic.

My attempt is:

var thingsToMock = new List<(string PartitionKey, string RowKey, string Value)>() {
    ("maxCount", "maxCount", "0"),
    ("maxCount", "xyz", "1000"),
    ("maxCount", "abc", "2000")
};
var tableClientMock = new Mock<TableClient>();
foreach (var thingToMock in thingsToMock)
{
    var returnPage = Page<SomeClass>.FromValues(new List<SomeClass>
    {
        new SomeClass{ PartitionKey = thingToMock.PartitionKey, RowKey = thingToMock.RowKey, Value = thingToMock.Value }
    }, null, new Mock<Response>().Object);
    var returnPages = AsyncPageable<SomeClass>.FromPages(new[] { returnPage });
    Expression<Func<SomeClass, bool>> exp = (x) => x.PartitionKey == thingToMock.PartitionKey && x.RowKey == thingToMock.RowKey ? true : false;
    tableClientMock
        .Setup(i => i.QueryAsync<SomeClass>(It.Is<Expression<Func<SomeClass, bool>>>(expression => LambdaExpression.Equals(expression, exp)), null, null, default))
        .Returns(returnPages);
}

The issue is that the _tableClientMock doesn't seem to return what I expected when I call GetAsync("maxCount", "abc"). I'd expect with this call, it would pass in the same parameters to tableClient.QueryAsync() method, which in my Mock should return instance of SomeClass with value of 2000. But instead, it throw "Object reference not set to an instance of an object." error.

If I change the tableClientMock setup for QueryAsync to be the following, it somewhat works:

.Setup(i => i.QueryAsync<SomeClass>(It.IsAny<Expression<Func<SomeClass, bool>>>(), null, null, default))

But this will not achieve my objective, which is to be able to pass different parameters (partitionKey and rowKey) to get different result.

I'm using the following NuGet package:

• "Azure.Data.Tables" Version="12.7.1"
• "moq" Version="4.14.1"

So I ended up with the following test. This is lil bit different than the example I provided in the question. But generally it's the same concept on how to mock Azure TableClient

The unit test

// Arrange
var expectedItems = new List<ItemEntity>
{
    new ItemEntity { RowKey = $"0001", PartitionKey = "abc" },
    new ItemEntity { RowKey = $"0002", PartitionKey = "abc" },
}
var tableClient = new Mock<TableClient>();
tableClient
    .Setup(c => c.QueryAsync<ItemEntity>(It.IsAny<string>(),
        It.Is<int>(a => a == 2),
        It.IsAny<IEnumerable<string>>(),
        It.IsAny<CancellationToken>()))
    .Returns(new MockAsyncPageable<ItemEntity>(expectedItems));
var repo = new Service(tableClient);

// Act
var result = await repo.GetItems("000", 2);

// Assert
result.Should().NotBeNull();
result.Any().Should().BeTrue();
result.Count().Should().Be(2);
tableClient.Verify(c => c.QueryAsync<ItemEntity>(It.IsAny<string>(),
    It.Is<int>(a => a == 2),
    It.IsAny<IEnumerable<string>>(),
    It.IsAny<CancellationToken>()), Times.Once);

The Service's GetItems method looks like this:

public async Task<IEnumerable<ItemEntity>> GetItems(string code, int maxPerPage)
{
    var maxPerPage = Math.Clamp(maxPerPage, 1, 1000);
    var filter = $"(RowKey eq '{partialNdc}') and (PartitionKey eq 'abc')";
    return await _tableClient.QueryAsync<ItemEntity>(filter, maxPerPage);
}

And lastly the, MockAsyncPageable<>

internal class MockAsyncPageable<T> : AsyncPageable<T>
{
    private readonly List<T> _items;

    public MockAsyncPageable(List<T> items)
    {
        _items = items;
    }

    public override IAsyncEnumerable<Page<T>> AsPages(string continuationToken = null, int? pageSizeHint = null)
    {
        var page = Page<T>.FromValues(_items, null, new Mock<Response>().Object);
        return new[] { page }.ToAsyncEnumerable();
    }
}

I believe the key is to setup the Mock (.Setup()) for the method I'm using, specifically, pass in all the parameters (including optional parameters) even though I don't use it in my GetItems() method.

For example, the .QueryAsync in my GetItems() method used the following method signature of TableClient:

public virtual AsyncPageable<T> QueryAsync<T>(string filter = null, int? maxPerPage = null, IEnumerable<string> select = null, CancellationToken cancellationToken = default(CancellationToken)) where T : class, ITableEntity { }

So when I setup the Mock, I pass in all the parameters.