diff --git a/app/main/views/two_factor.py b/app/main/views/two_factor.py index 0be1239fc..a6af6c094 100644 --- a/app/main/views/two_factor.py +++ b/app/main/views/two_factor.py @@ -70,15 +70,16 @@ def two_factor(): return user_api_client.check_verify_code(user_id, code, "sms") form = TwoFactorForm(_check_code) + redirect_url = request.args.get('next') if form.validate_on_submit(): if is_less_than_90_days_ago(user.email_access_validated_at): return log_in_user(user_id) else: - user_api_client.send_verify_code(user.id, 'email', None, request.args.get('next')) - return redirect(url_for('.revalidate_email_sent')) + user_api_client.send_verify_code(user.id, 'email', None, redirect_url) + return redirect(url_for('.revalidate_email_sent', next=redirect_url)) - return render_template('views/two-factor.html', form=form) + return render_template('views/two-factor.html', form=form, redirect_url=redirect_url) @main.route('/re-validate-email', methods=['GET']) diff --git a/app/templates/views/two-factor.html b/app/templates/views/two-factor.html index 8a7cfa2f7..d0a252f14 100644 --- a/app/templates/views/two-factor.html +++ b/app/templates/views/two-factor.html @@ -22,7 +22,7 @@ }) }} {{ page_footer( "Continue", - secondary_link=url_for('main.check_and_resend_text_code'), + secondary_link=url_for('main.check_and_resend_text_code', next=redirect_url), secondary_link_text='Not received a text message?' ) }} {% endcall %} diff --git a/tests/app/main/views/test_two_factor.py b/tests/app/main/views/test_two_factor.py index 5fc08934f..15f6c5568 100644 --- a/tests/app/main/views/test_two_factor.py +++ b/tests/app/main/views/test_two_factor.py @@ -37,11 +37,16 @@ def test_two_factor_email_sent_page( assert resend_email_link['href'] == url_for('main.email_not_received', next=redirect_url) +@pytest.mark.parametrize('redirect_url', [ + None, + 'blob', +]) def test_should_render_two_factor_page( client, api_user_active, mock_get_user_by_email, - mocker + mocker, + redirect_url ): # TODO this lives here until we work out how to # reassign the session after it is lost mid register process @@ -50,7 +55,7 @@ def test_should_render_two_factor_page( 'id': api_user_active['id'], 'email': api_user_active['email_address']} mocker.patch('app.user_api_client.get_user', return_value=api_user_active) - response = client.get(url_for('main.two_factor')) + response = client.get(url_for('main.two_factor', next=redirect_url)) assert response.status_code == 200 page = BeautifulSoup(response.data.decode('utf-8'), 'html.parser') assert page.select_one('main p').text.strip() == ( @@ -62,6 +67,10 @@ def test_should_render_two_factor_page( assert page.select_one('input')['type'] == 'tel' assert page.select_one('input')['pattern'] == '[0-9]*' + assert page.select_one( + 'a:contains("Not received a text message?")' + )['href'] == url_for('main.check_and_resend_text_code', next=redirect_url) + @freeze_time('2020-01-27T12:00:00') def test_should_login_user_and_should_redirect_to_next_url( @@ -104,11 +113,15 @@ def test_should_send_email_and_redirect_to_info_page_if_user_needs_to_revalidate session['user_details'] = { 'id': api_user_active['id'], 'email': api_user_active['email_address']} - response = client.post(url_for('main.two_factor', next='/services/{}'.format(SERVICE_ONE_ID)), + response = client.post(url_for('main.two_factor', next=f'/services/{SERVICE_ONE_ID}'), data={'sms_code': '12345'}) assert response.status_code == 302 - assert response.location == url_for('main.revalidate_email_sent', _external=True) + assert response.location == url_for( + 'main.revalidate_email_sent', + _external=True, + next=f'/services/{SERVICE_ONE_ID}' + ) mock_send_verify_code.assert_called_with(api_user_active['id'], 'email', None, mocker.ANY)