Zenn
🐍

いつものpytestに1行追加してメモリリークを見つける

2024/06/06に公開
Python
Test
pytest
memray
tech

はじめに

Pythonの単体テスト実施時に、各機能のメモリ使用量やメモリリークを計測するライブラリを紹介します。Pythonで書かれたコードに対してメモリ計測を行うライブラリは複数ありますが、その中で最もモダンなライブラリとしてmemrayがあります。

https://bloomberg.github.io/memray/

今回はmemrayからpytest用に切り出されたpytest-memrayというライブラリを紹介します。

https://pytest-memray.readthedocs.io/en/latest/

pytest-memray

インストール

インストールはpip installで可能です。

pip install pytest-memray

メモリ使用量に関するテスト

まずデモとして以下のようなテストコードを作成します。

import pytest


def test_foobar():
    number_list = []
    for i in range(100000):
        number_list.append(i)

    assert isinstance(number_list, list)
    assert len(number_list) == 100000

このようなテストコードを実装することは少ないと思いますが、今回はデモ用として用意しています。上記のテストでは、number_listというリストに大量のint型整数が格納されるため、多くのメモリが割り当てられます。つまり、テスト自体は通るものの、メモリ使用量の増加には気付けない状態になってしまいます。

そこでpytest-memrayを用いて、以下のように実装します。

+ @pytest.mark.limit_memory("1 MB")
def test_foobar():
    number_list = []
    for i in range(100000):
        number_list.append(i)

    assert isinstance(number_list, list)
    assert len(number_list) == 100000

メモリ使用量を測定し、上記の例では1MB以上メモリを使用した場合にアラートが発生します。テストを実行する際は--memrayのコマンドラインオプションを付ける必要があります。

pytest --memray

以下実行結果です。

====================================================================================================== MEMRAY REPORT =======================================================================================================
Allocation results for tests/test_sample.py::test_foobar at the high watermark

         📦 Total memory allocated: 3.8MiB
         📏 Total allocations: 3
         📊 Histogram of allocation sizes: |█ |
         🥇 Biggest allocating functions:
                - test_foobar:/workspace/tests/test_sample.py:7 -> 3.0MiB
                - test_foobar:/workspace/tests/test_sample.py:8 -> 782.2KiB


================================================================================================= short test summary info ==================================================================================================
MEMORY PROBLEMS tests/test_sample.py::test_foobar
==================================================================================================== 1 failed in 0.04s =====================================================================================================

合計3.8MBのメモリを使用しており、1MBを大きく超えるためテストは失敗に終わっています。

メモリリークに関するテスト

メモリリークについても計測可能で、その際はpytest.mark.limit_leaksを使用します。

+ @pytest.mark.limit_leaks("1 MB")
def test_foobar():
    number_list = []
    for i in range(100000):
        number_list.append(i)

    assert isinstance(number_list, list)
    assert len(number_list) == 100000

まとめ

pytest-memrayを用いることで、通常の機能テストに加えて、メモリ使用量やメモリリークのテストも行うことができます。また実装方法は非常にシンプルで@pytest.mark.limit_memoryやpytest.mark.limit_memoryを付けるだけです。

以上です！

Discussion