この記事はエムスリーAdvent Calendar 2023の18日目の記事です。 AI・機械学習チームの高田です。
AI・機械学習チームではデータパイプラインを構築する機会が多く、パイプラインの中でpandasを活用しています。 今回はpandasのSeries型を扱う関数の単体テストにMagicMockを使った際にハマったポイントを紹介したいと思います。
Series.mapの活用例
データの前処理工程ではDataFrameやSeriesの値を利用して別のSeriesを作成することが多くあります。 例えばアンケートの自由記述回答のうち、一定の文字数以上のデータのみ扱うために、テキストデータのSeriesから文字列長のSeriesを作成する場合が考えられます。 SeriesからSeriesの作成にはmapやapplyをよく使いますが、今回はmapを使用してみます。
サンプル実装です。
import pandas as pd from typing import Callable def add_column_from_text(df: pd.DataFrame, column_name: str, mapper: Callable): df[column_name] = df['text'].map(mapper) return df def calc_text_length(x): return len(x) df = pd.DataFrame([['おはようございます'], ['こんにちは']], columns=['text']) new_df = add_column_from_text(df, 'length', calc_text_length) sufficient_length_df = new_df[new_df['length'] > 5] print(sufficient_length_df)
lengthカラムに文字列長を格納できています。
| text | length |
|---|---|
| おはようございます | 9 |
単体テスト
では、add_column_from_textの単体テストを書きます。
import unittest class TestAddColumnFromText(unittest.TestCase): def test_add_column_from_text(self): """ - 返り値が期待通りである - mapperが意図通りの引数で呼ばれている ことを保証する """ expected = pd.DataFrame([ dict(text='おはようございます', length=9), dict(text='こんにちは', length=5), ]) mock_func = MagicMock() mock_func.side_effect = calc_text_length actual = add_column_from_text(df, 'length', mock_func) pd.testing.assert_frame_equal(expected, actual) mock_func.assert_any_call('おはようございます') mock_func.assert_any_call('こんにちは')
実行結果
===================================================================== FAIL: test_add_column_from_text (__main__.TestAddColumnFromText) mapperが意図通りの引数で呼ばれていることを保証する ---------------------------------------------------------------------- Traceback (most recent call last): ... 略 ... AssertionError: Attributes of DataFrame.iloc[:, 1] (column name="length") are different Attribute "dtype" are different [left]: int64 [right]: float64 ---------------------------------------------------------------------- Ran 1 test in 0.023s FAILED (failures=1)
actualを表示してみると lengthが全てNaNになっています。
| text | length |
|---|---|
| おはようございます | NaN |
| こんにちは | NaN |
何が起きているのか
mock自体はMagicMockで正しく定義できているようです。
>>> mock_func = MagicMock() >>> mock_func.side_effect = calc_text_length >>> mock_func('こんにちは') 5
となると、pandas.Series.mapの内部で意図しない動作になっていそうです。 pandas.Series.mapのAPI referenceを確認すると
arg: function, collections.abc.Mapping subclass or Series
mapの第一引数は関数、 dictのサブクラス、 Series型を受け付けます。
When arg is a dictionary, values in Series that are not in the dictionary (as keys) are converted to NaN. However, if the dictionary is a dict subclass that defines missing (i.e. provides a method for default values), then this default is used rather than NaN.
その第一引数がdict型で、defaultdictのようにkeyが存在しないときのデフォルト値を定義していない場合にNaNで埋める挙動をするという注意書きがあります。
どうやらmock_funcがdict型として認識されて関数として実行されていないのではないか？という仮説を立てました。
そこでpandas.Series.mapの内部実装を追っていきます。 関数コールを追っていくとmap_arrayの中に
if is_dict_like(mapper):
という記述があります。mapperはmock_funcに相当します。ここでdictライクかそれ以外かで処理が分岐していました。
is_dict_like関数の実装を確認すると以下3つのattributeを持つオブジェクトをdictライクと判定しています。
__getitem__keys__contains__
>>> [hasattr(mock_func, attr) for attr in ("__getitem__", "keys", "__contains__")] [True, True, True]
mock_funcはこれら3つのattributeを持っていたためにdicライクとして判定され、関数として実行されていないことが原因だということが判明しました。
MagicMockは通常のMockの機能に加え、__getitem__や__contains__も含む多くのmagic methodがサポートされています。
今回はmagic methodのサポートが仇となり、関数として判定される前にdict型として判定されてしまったことがバグの原因でした。
修正結果
そこで、単体テストを以下のように修正しました。
import unittest
class TestAddColumnFromText(unittest.TestCase):
def test_add_column_from_text(self):
"""
- 返り値が期待通りである
- mapperが意図通りの引数で呼ばれている
ことを保証する
"""
expected = pd.DataFrame([
dict(text='おはようございます', length=9),
dict(text='こんにちは', length=5),
])
- mock_func = MagicMock()
+ mock_func = Mock()
mock_func.side_effect = calc_text_length
actual = add_column_from_text(df, 'length', mock_func)
pd.testing.assert_frame_equal(expected, actual)
mock_func.assert_any_call('おはようございます')
mock_func.assert_any_call('こんにちは')
実行結果
---------------------------------------------------------------------- Ran 1 test in 0.008s OK
テストが通りました。
おわりに
今までは便利なMockだなという軽い気持ちでMagicMockを使っていましたが、MagicMockを使用するべきではない事例に初めて遭遇しました。
個人的にはOSSの内部実装を地道に読み解くよい機会になり、勉強になりました。
We are hiring!!
AI・機械学習チームでは絶賛エンジニアを募集中です！ 作って終わりじゃないテスタブルな実装にこだわりを持ってプロダクト開発に取り組みたい方は是非カジュアル面談等ご応募ください！
