production環境のDatastoreを使って手元でUnitTestを実行する

前にremote_api_shell.pyを使って対話実行環境でproductionのDatastoreに接続する方法を書いた。
で、もちろん次は自動テストをproductionのDatastoreでやりたいよね？ということで、そのための方法を書いたのでご紹介。Javaとかだとproduction環境内で直接テストを実行するKotoriWebJUnitRunner - kotori - a junit runner for GAE production - Kotori project is to know what I can do and what I don't know. - Google Project Hostingとかあるが、自分が書いたものはとりあえずコード自体は手元での実行で、すべてのDatastore RPCをremote_api_shellが使っているremote_api_stubを経由して走らせている。

(4/8追記：トランザクションはクライアント側で処理されてました。こちらを参照のこと。)

テスト用ベースクラスの準備

Python SDKでUnitTestを走らせるためには、普通に手元のDatastoreモドキに繋いで実行するときでも、環境の調整を自分でやる必要がある。nose-gae - nose plugin for working with google app engine applications - Google Project Hostingとか使えばラクらしいけど、自分はあまりフレームワークをいくつも組合せるのが嫌い*1なので、ちょっとぐぐって見付かったGoogle App Engine/Python で単体テスト - presentを参考にベースクラスを作ってます。

で、このコードを見るとGAETestBase.setUp()内でdatastore_file_stubなんかの初期化をやっているのがわかるので、あとはこれをremote_api_stubに置き換えてやればあら不思議productionのDatastoreに接続成功、というだけ。
ただし見境なくproductionに繋げてもちょっとマズいとかあれこれあって、実現したかった要求は以下の通り。

大量にテストコード書きたいけどTestCaseごとに繋ぎに行くのをproductionとdevで分けたいよね
productionに繋ぎに行くのにuser/pass必要だけどそれはコードに書きたくないよね
- TestCaseごとにuser/pass入力求められたらキレるから、どっかで一度入力したら使い回したいよね
productionのDatastoreの中身を汚すのはちょっとアレだから、TestCaseごとにクリーンアップの処理は決めたいよね

で、上記の要望を満たせるTestCaseクラス(のベースクラス)を以下のように作ってみた。

#!/usr/bin/env python2.5
# -*- coding:utf-8 -*-

import os, sys, functools, getpass

GAE_HOME = '/home/tagomoris/google_appengine'
PROJECT_HOME = '/home/tagomoris/tagomoris-test/trunk'
EXTRA_PATHS = [
    GAE_HOME,
    PROJECT_HOME,
    os.path.join(GAE_HOME, 'google', 'appengine', 'api'),
    os.path.join(GAE_HOME, 'google', 'appengine', 'ext'),
    os.path.join(GAE_HOME, 'lib', 'yaml', 'lib'),
    os.path.join(GAE_HOME, 'lib', 'webob'),
]
sys.path = EXTRA_PATHS + sys.path

import unittest

from google.appengine.api import apiproxy_stub_map
from google.appengine.runtime import apiproxy_errors

from google.appengine.api import datastore
from google.appengine.api import datastore_errors
from google.appengine.ext import db, search
from google.appengine.ext.db import polymodel

from google.appengine.ext.remote_api import remote_api_stub
from google.appengine.api import datastore_file_stub
from google.appengine.api import mail_stub
from google.appengine.api import urlfetch_stub
from google.appengine.api import user_service_stub

from google.appengine.api import users

APP_ID = u'tagomoris-test'
AUTH_DOMAIN = 'gmail.com'
LOGGED_IN_USER = 'test@example.com'

os.environ['APPLICATION_ID'] = APP_ID

def auth_func():
    return (raw_input('Email: '), getpass.getpass('Password: '))

class GAETestBase(unittest.TestCase):
    auth_pair = None

    def is_remote(self):
        """
        GAETestBaseを継承したテストケース用クラスで use_remote_stubs() が定義されて、かつTrueを返す
        場合にはremote_api_stub経由ですべてのRPCをproduction環境に飛ばす
        """
        return (getattr(self, 'use_remote_stubs', lambda: False))()
    
    def setUp(self):
        if self.is_remote():
            if GAETestBase.auth_pair is None:
                GAETestBase.auth_pair = auth_func()
            remote_api_stub.ConfigureRemoteApi(APP_ID, '/remote_api',
                                               lambda: GAETestBase.auth_pair)
            remote_api_stub.MaybeInvokeAuthentication()
        else:
            apiproxy_stub_map.apiproxy = apiproxy_stub_map.APIProxyStubMap()

            # dummy local datastore
            stub = datastore_file_stub.DatastoreFileStub(APP_ID, '/dev/null', '/dev/null')
            apiproxy_stub_map.apiproxy.RegisterStub('datastore_v3', stub)
            
            apiproxy_stub_map.apiproxy.RegisterStub('user', user_service_stub.UserServiceStub())
            os.environ['AUTH_DOMAIN'] = AUTH_DOMAIN
            os.environ['USER_EMAIL'] = LOGGED_IN_USER
            apiproxy_stub_map.apiproxy.RegisterStub('urlfetch', urlfetch_stub.URLFetchServiceStub())
            apiproxy_stub_map.apiproxy.RegisterStub('mail', mail_stub.MailServiceStub()) 

    def tearDown(self):
        """
        GAETestBaseを継承したテストケース用クラスで cleanup_remote_stubs() が定義されて
        いる場合には呼び出す
        """
        if self.is_remote() and getattr(self, 'cleanup_remote_stubs', False):
            self.cleanup_remote_stubs()

これを gae_test_base.py に書いたとすると、実際にテストを書きたい hoge_test.py では次のように書く。

#!/usr/bin/env python2.5
# -*- coding:utf-8 -*-

from gae_test_base import *
from test_models import *

from google.appengine.ext import db

class HogeTest(GAETestBase):
    def use_remote_stubs(self):
        return True

    def cleanup_remote_stubs(self):
        """
        ここでは試験用のKind XEntity を全部 delete() している
        """
        q = XEntity.all(keys_only=True)
        while True:
            l = q.fetch(500)
            if len(l) < 1:
                break
            db.delete(l)

    def test_hogehoge(self):
        """ ここで実行したいテストの内容を記述"""
        x = XEntity(key_name="xxx1")
        x.put()

if __name__ == '__main__':
    unittest.main()

もちろん実行の前提として、自分のアプリケーション(production)に対して remote_api_shell.py を通じたアクセスが可能な状態が必要です。
実行すると一度だけEmailアドレスとパスワードを聞かれ、それを使って自分のアプリケーションに接続して処理を行う。当然すべての処理をHTTPを通じてRPCするから処理速度はものすごく遅いけど、まあ、しょうがないでしょう、ということで。

速度を考えるならproduction内で全部実行した方がいいと思うけど、自分の用途はちょっとdev/productionを切り替えながら走らせたかったので、とりあえずこれでいいかなと思ってます。
KotoriWebJUnitRunnerみたいなののPython版は全世界に期待されてるので誰か！ｗ

*1:GAE SDKは1.3.2になったけどそこじゃあまだこのフレームワークは動かないんでアップデートは待って！とか言われたら憤死しそう