I’ve been working on a project for my company https://gitlab.com/paessler-labs/prtg-pyprobe and it’s using asyncio library for Python quite heavily. It’s a great framework and so far we’ve had very little issues implementing it but I recently implemented some new code that I wanted to mock for unit testing and it was a bit more challenging than I thought it would be.
Here is the code that I wanted to test
start = time.time()
async with aioboto3.resource(
"s3", aws_access_key_id=task_data["aws_access_key"], aws_secret_access_key=task_data["aws_secret_key"]
) as s3:
all_buckets = s3.buckets.all()
i = 0
async for _ in all_buckets:
i += 1
s3_bucket_data.add_channel(
name="Total Buckets", mode="integer", kind="Custom", customunit="buckets", value=i
)
s3_bucket_data.message = f"Your AWS account has {i} buckets."
end = (time.time() - start) * 1000
s3_bucket_data.add_channel(name="Total Query Time", mode="float", kind="TimeResponse", value=end)
And here is what the unit test ended up looking like
@pytest.mark.asyncio
class TestS3TotalWork:
@asynctest.patch("aioboto3.resource")
async def test_sensor_s3_total(self, aioboto_mock, s3_total_sensor):
buckets = asynctest.MagicMock()
buckets.__aiter__.return_value = ["bucket1", "bucket2", "bucket3"]
aioboto_mock.return_value.__aenter__.return_value.buckets.all.return_value = buckets
s3_total_queue = asyncio.Queue()
await s3_total_sensor.work(task_data=task_data(), q=s3_total_queue)
queue_result = await s3_total_queue.get()
aioboto_mock.assert_called_once_with("s3", aws_access_key_id="1123124", aws_secret_access_key="jkh2089")
assert queue_result["message"] == "Your AWS account has 3 buckets."
assert {
"customunit": "buckets",
"kind": "Custom",
"mode": "integer",
"name": "Total Buckets",
"value": 3,
} in queue_result["channel"]
To be able to test this, I ended up using the https://pypi.org/project/asynctest/ which was really helpful for mocking async iterables and context managers.
The trick for me to figure out was what was returning what when for the call
all_buckets = s3.buckets.all()
The way that I tried to think of it is as such:
@asynctest.patch("aioboto3.resource")
...
aioboto_mock.return_value.__aenter__.return_value.buckets.all.return_value = buckets
‘aioboto_mock’ mocks the library and it’s attribute ‘resource’
The ‘return_value’ of that is the ‘s3’ context manager
Then we step into the context which is the ‘__aenter__’ method
Then we need another ‘return_value’ since this is where the context is returning the results from..
The method buckets.all()’s return_value.
buckets = asynctest.MagicMock()
buckets.__aiter__.return_value = ["bucket1", "bucket2", "bucket3"]
The buckets.all() method also returns an async iterable and so to patch the result of the s3.buckets.all() method we also have to set this to be a MagicMock that returns values for the __aiter__ method of async iterables..
Totally simple right :D. Hope this helps anyone else trying to understand how to mock async functions!
opendiscourse not only about monitoring in the technical sense but also about the human side of monitoring which isn’t something that gets a lot of press.
Recent Comments