[Python] 環境変数のテストで役立つreload

TL; DR

configやsettingsでの環境変数を動的に変えるテストで importlib.reload が便利。jupyterでの検証時も使える。

状況

settingsなどに環境変数からの読み込みを集めるということをすることもあるかと思います。

settings.py

import os
API_KEY = os.getenv("API_KEY") 

何らかの事情で新しい環境変数を入れて上書きをしたいというのを、旧環境も保ったままやりたくなったとしましょう¹。

settings.py

import os
API_KEY = os.getenv("API_KEY") 
API_KEY_ALT = os.getenv("API_KEY_ALT") 
if API_KEY_ALT is not None:
    API_KEY = API_KEY_ALT

何だかsettingsにロジックが入ってしまいました。ロジックが入ってしまったからにはテストを書かなければいけません²。

テストを書こう

環境変数のテストはos.environをpatchすればよさそうです。

test_settings.py

import os
from unittest import TestCase
from unittest import mock


class MyTestCase(TestCase):
    def test_1(self):
        env = {}
        with mock.patch.dict(os.environ, env):
            import settings
        self.assertIsNone(settings.API_KEY)

環境変数として何もない状況。これの結果はAPI_KEYはNoneになりそうです。これは通ります。

test_settings.py

import os
from unittest import TestCase
from unittest import mock


class MyTestCase(TestCase):
    def test_2(self):
        env = {"API_KEY": "api"}
        with mock.patch.dict(os.environ, env):
            import settings
        self.assertEqual(settings.API_KEY, "api")

環境変数に"api"を入れました。当然、このテストも通ります。ですが、

test_settings.py

import os
from unittest import TestCase
from unittest import mock


class MyTestCase(TestCase):
    def test_1(self):
        env = {}
        with mock.patch.dict(os.environ, env):
            import settings
        self.assertIsNone(settings.API_KEY)

    def test_2(self):
        env = {"API_KEY": "api"}
        with mock.patch.dict(os.environ, env):
            import settings
        self.assertEqual(settings.API_KEY, "api")

こうすると2番目のテストが通らなくなってしまいました。

インポートが邪魔をする

これの理由は2回目のインポートが1回目のインポートに影響を受けていることによります。test結果を見ていると2回目のAPI_KEYもNoneになっているというのがわかります。
これを解決するのがimportlib.reloadです。この関数

以前にインポートされた module をリロードします。引数はモジュールオブジェクトでなければならず、したがってそれ以前に必ずインポートに成功していなければなりません。この関数は、モジュールのソースファイルを外部エディタで編集していて Python インタープリタから離れることなく新しいバージョンを試したい際に便利です。戻り値はモジュールオブジェクトです。

ということで呼びたいモジュールをコードレベルで読み込んでくれます。これを利用すると

test_settings.py

from importlib import reload
import os
from unittest import TestCase
from unittest import mock


class MyTestCase(TestCase):
    def test_1(self):
        env = {}
        with mock.patch.dict(os.environ, env):
            import settings
        self.assertIsNone(settings.API_KEY)

    def test_2(self):
        env = {"API_KEY": "api"}
        with mock.patch.dict(os.environ, env):
            import settings
            reload(settings)
        self.assertEqual(settings.API_KEY, "api")

これで通るようになります。これでやっと書きたかったテストが書けるようになりました。

他の使い方

このreloadですが、この記事によるとjupyterで動作検証しているときにも便利そうです。関数を書くたびにカーネルを上げ直す日々にさよならできますね。

1. そんな横着するな。という状況。 ↩

2. ちゃんと先に書きましょう。 ↩