36 lines
1.0 KiB
Python
36 lines
1.0 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
import flask
|
||
|
|
|
||
|
|
|
||
|
|
class FlaskExtraTest(unittest.TestCase):
|
||
|
|
def test_vary_xhr(self):
|
||
|
|
import pillar.flask_extra
|
||
|
|
|
||
|
|
class TestApp(flask.Flask):
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
self.add_url_rule('/must-vary', 'must-vary', self.must_vary)
|
||
|
|
self.add_url_rule('/no-vary', 'no-vary', self.no_vary)
|
||
|
|
|
||
|
|
@pillar.flask_extra.vary_xhr()
|
||
|
|
def must_vary(self):
|
||
|
|
return 'yay'
|
||
|
|
|
||
|
|
def no_vary(self):
|
||
|
|
return 'nah', 201
|
||
|
|
|
||
|
|
app = TestApp(__name__)
|
||
|
|
client = app.test_client()
|
||
|
|
|
||
|
|
resp = client.get('/must-vary')
|
||
|
|
self.assertEqual(200, resp.status_code)
|
||
|
|
self.assertEqual('X-Requested-With', resp.headers['Vary'])
|
||
|
|
self.assertEqual('yay', resp.data.decode())
|
||
|
|
|
||
|
|
resp = client.get('/no-vary')
|
||
|
|
self.assertEqual(201, resp.status_code)
|
||
|
|
self.assertNotIn('Vary', resp.headers)
|
||
|
|
self.assertEqual('nah', resp.data.decode())
|