时间戳

首页 > 解决方案 > 使用 csrf 令牌测试表单提交

flask - 使用 csrf 令牌测试表单提交

问题描述

创建新帐户时，我将用户重定向到他的个人信息摘要。但如果不返回
则不会重定向。validate_on_submitTrue

如果CSRF_TOKEN匹配则验证。

如何csrf_token使用 pytest 在测试环境中创建和附加一个？

这是我的测试：

@pytest.fixture
def client():
    app = create_app({"TESTING": True})

    with app.test_client() as cl:
        yield cl

def test_new_first(client):
    rv = client.post("/signup", data={
        "username": "John Doe",
        "email": "johndoe@gmail.com",
        "password": "11111111",
    }, follow_redirects=True)

    assert b"Summary" in rv.data

问题是：即使follow_redirects=True它停留在注册页面上并且“摘要”不在rv.data.
我应该csrf_token把data.

我怎样才能生成一个csrf_token？

标签： flaskpytestflask-wtforms

解决方案

我可以根据这个 gist的工作来测试包含 CSRF 令牌的表单。为方便起见，我添加了一个通用post_csrf方法：

# configure test_client to handle CSRF tokens properly
# according to https://gist.github.com/singingwolfboy/2fca1de64950d5dfed72

# Flask's assumptions about an incoming request don't quite match up with
# what the test client provides in terms of manipulating cookies, and the
# CSRF system depends on cookies working correctly. This little class is a
# fake request that forwards along requests to the test client for setting
# cookies.
class RequestShim(object):
    """
    A fake request that proxies cookie-related methods to a Flask test client.
    """
    def __init__(self, client):
        self.client = client
        self.vary = set({})

    def set_cookie(self, key, value='', *args, **kwargs):
        """Set the cookie on the Flask test client."""
        server_name = flask.current_app.config['SERVER_NAME'] or 'localhost'
        return self.client.set_cookie(
            server_name, key=key, value=value, *args, **kwargs
        )

    def delete_cookie(self, key, *args, **kwargs):
        """Delete the cookie on the Flask test client."""
        server_name = flask.current_app.config['SERVER_NAME'] or 'localhost'
        return self.client.delete_cookie(
            server_name, key=key, *args, **kwargs
        )


# We're going to extend Flask's built-in test client class, so that it knows
# how to look up CSRF tokens for you!
class FlaskClient(BaseFlaskClient):
    @property
    def csrf_token(self):
        # First, we'll wrap our request shim around the test client, so that
        # it will work correctly when Flask asks it to set a cookie.
        request = RequestShim(self)
        # Next, we need to look up any cookies that might already exist on
        # this test client, such as the secure cookie that powers `flask.session`,
        # and make a test request context that has those cookies in it.
        environ_overrides = {}
        self.cookie_jar.inject_wsgi(environ_overrides)
        with flask.current_app.test_request_context(
                '/auth/login', environ_overrides=environ_overrides,
            ):
            # Now, we call Flask-WTF's method of generating a CSRF token...
            csrf_token = generate_csrf()
            # ...which also sets a value in `flask.session`, so we need to
            # ask Flask to save that value to the cookie jar in the test
            # client. This is where we actually use that request shim we made!
            flask.current_app.session_interface.save_session(flask.current_app, flask.session, request)
            # And finally, return that CSRF token we got from Flask-WTF.
            return csrf_token

    # Feel free to define other methods on this test client. You can even
    # use the `csrf_token` property we just defined, like we're doing here!
    def login(self, username='test', password='test'):
        # use post_csrf instead of code of linked gist
        return self.post_csrf('/auth/login', username=username, password=password, remember_me=False)

    def logout(self):
        return self.get('/auth/logout', follow_redirects=True)

    # generic post with csrf_token to test all form submissions of my flask app
    def post_csrf(self, url, **kwargs):
        data = kwargs
        data['csrf_token'] = self.csrf_token

        return self.post(url, data=data, follow_redirects=True)


@pytest.fixture
def app():
    """Create and configure a new app instance for each test."""
    app = create_app(TestConfig)
    app.test_client_class = FlaskClient
    
    # rest omitted for brevity


@pytest.fixture
def client(app):
    """A test client for the app."""
    return app.test_client()

然后测试如下所示：

def test_remove_nonexisting_user_fails_gracefully(client):
    u = User(username='alice', password='pass')
    db.session.add(u)
    db.session.commit()

    client.login()
    response = client.post_csrf('/users/delete/bob')
    assert response.status_code == 200
    assert len(user.query.all()) == 1
    assert b'user bob not found' in response.data

推荐阅读