def test_get_db(app):
with app.app_context():
db = get_db()
assert db is get_db()
with pytest.raises(sqlite3.ProgrammingError) as e:
db.execute('SELECT 1')
assert 'closed' in str(e)
def app():
db_fd, db_path = tempfile.mkstemp()
app = create_app(test_config={'TESTING': True, 'DATABASE': db_path})
with app.app_context():
init_db()
get_db().executescript(_data_sql)
yield app
os.close(db_fd)
os.unlink(db_path)
def login():
username = request.form['username']
password = request.form['password']
db = get_db()
error_msg = None
if not username:
error_msg = 'Username is required.'
elif not password:
error_msg = 'Password is required.'
else:
user = db.execute('SELECT * FROM user WHERE username = ?',
(username, )).fetchone()
if user is None or not check_password_hash(user['password'], password):
error_msg = 'Incorrect username or password.'
if error_msg is None:
session.clear()
session['user_id'] = user['id']
flash('Login success!', category='info')
return redirect(url_for('index'))
flash(error_msg, category='error')
return render_template('index.html')
def register():
username = request.form['username']
password = request.form['password']
re_password = request.form['re_password']
db = get_db()
error_msg = None
if not username:
error_msg = 'Username is required.'
elif not password or not re_password:
error_msg = 'Password is required.'
elif password != re_password:
error_msg = 'Password is not the same.'
elif db.execute('SELECT id FROM user where username = ?',
(username, )).fetchone() is not None:
error_msg = 'Username {} has alread been registered.'.format(username)
if error_msg is None:
db.execute('INSERT INTO user (username, password) VALUES (?, ?)',
(username, generate_password_hash(password)))
db.commit()
user = db.execute('SELECT * FROM user WHERE username = ?',
(username, )).fetchone()
session.clear()
session['user_id'] = user['id']
flash('Register success!', category='info')
return redirect(url_for('index'))
flash(error_msg, category='error')
return render_template('index.html')
def index():
if request.method == 'POST':
long_URL = request.form['long_URL']
err_msg = None
shortURL = None
if not long_URL or long_URL == "":
err_msg = "Please input URL."
else:
db = get_db()
url = db.execute("SELECT * FROM url WHERE long_URL = ?",
(long_URL, )).fetchone()
if url:
shortURL = shorten_URL(url['id'])
else:
db.execute(
"INSERT INTO url (long_URL, owner_id) values (?, ?)",
(long_URL, 1))
db.commit()
url = db.execute("SELECT * FROM url WHERE long_URL = ?",
(long_URL, )).fetchone()
shortURL = shorten_URL(['id'])
return render_template('index.html', short_URL=shortURL)
flash(err_msg, 'error')
return render_template('index.html', short_URL=None)
def load_user_info():
"""
Load user information if user has logged in before handling every requests.
"""
user_id = session.get('user_id')
if user_id is None:
g.user = None
else:
g.user = get_db().execute('SELECT * FROM user WHERE id = ?',
(user_id, )).fetchone()
def test_register(app, client):
response = client.post('/auth/register',
data={
'username': '******',
'password': '******',
're_password': '******'
})
assert response.headers['Location'] == 'http://localhost/'
with app.app_context():
assert get_db().execute(
"SELECT * from user WHERE username = '******'").fetchone() is not None