テスト時と本番環境で処理を分けたい

テスト時とかにどうしても処理分けたかった。sys.argvを使って実行時引数のファイル名から判断しようかな。

import sys

if 'test_loader.py' in sys.argv:
    db = load_database('test')
else:
    db = load_database('production')

実行時

python test_loader.py

でも、もし test_loaderを他のスクリプトから呼びだす必要がある場合、この方法はだめ。。

※ 2014/11/12 追記
コメントで頂いた環境変数を設定する場合は os.environ を直接操作したほうが良いとのこと。
15.1. os — 雑多なオペレーティングシステムインタフェース Python 2.7ja1 documentation

test_loader.py:

import os
if not os.getenv('DB'):
    os.environ['DB'] = 'test'