From a6e1210162681b16d437218b4e9df46143ff0aee Mon Sep 17 00:00:00 2001
From: Rushabh Mehta <rmehta@gmail.com>
Date: Tue, 8 Aug 2017 14:42:02 +0530
Subject: [PATCH] [fix] tests

---
 frappe/tests/test_twofactor.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/frappe/tests/test_twofactor.py b/frappe/tests/test_twofactor.py
index 900e617360..e993b2d517 100644
--- a/frappe/tests/test_twofactor.py
+++ b/frappe/tests/test_twofactor.py
@@ -5,27 +5,25 @@ from __future__ import unicode_literals
 import unittest, frappe, pyotp
 from werkzeug.wrappers import Request
 from werkzeug.test import EnvironBuilder
-from frappe.auth import LoginManager, HTTPRequest
-from frappe.twofactor import *
-import time
-
+from frappe.auth import HTTPRequest
+from frappe.twofactor import (should_run_2fa, authenticate_for_2factor, get_cached_user_pass,
+	two_factor_is_enabled_for_, confirm_otp_token, get_otpsecret_for_, get_verification_obj,
+	render_string_template)
 
+import time
 
 class TestTwoFactor(unittest.TestCase):
-
-
 	def setUp(self):
 		self.http_requests = create_http_request()
 		self.login_manager = frappe.local.login_manager
 		self.user = self.login_manager.user
 
 	def tearDown(self):
-		tmp_id = frappe.local.response['tmp_id']
 		frappe.local.response['verification'] = None
 		frappe.local.response['tmp_id'] = None
+		disable_2fa()
 		frappe.clear_cache(user=self.user)
 
-
 	def test_should_run_2fa(self):
 		'''Should return true if enabled.'''
 		toggle_2fa_all_role(state=True)
@@ -68,7 +66,7 @@ class TestTwoFactor(unittest.TestCase):
 		otp = 'wrongotp'
 		with self.assertRaises(frappe.AuthenticationError):
 			confirm_otp_token(self.login_manager,otp=otp,tmp_id=tmp_id)
-		otp = get_otp(self.user)		
+		otp = get_otp(self.user)
 		self.assertTrue(confirm_otp_token(self.login_manager,otp=otp,tmp_id=tmp_id))
 		if frappe.flags.tests_verbose:
 			print('Sleeping for 30secs to confirm token expires..')
@@ -106,12 +104,18 @@ def create_http_request():
 
 def enable_2fa():
 	'''Enable Two factor in system settings.'''
-	toggle_two_factor_auth(True)
 	system_settings = frappe.get_doc('System Settings')
+	system_settings.enable_two_factor_auth = 1
 	system_settings.two_factor_method = 'OTP App'
 	system_settings.save(ignore_permissions=True)
 	frappe.db.commit()
 
+def disable_2fa():
+	system_settings = frappe.get_doc('System Settings')
+	system_settings.enable_two_factor_auth = 0
+	system_settings.save(ignore_permissions=True)
+	frappe.db.commit()
+
 def toggle_2fa_all_role(state=None):
 	'''Enable or disable 2fa for 'all' role on the system.'''
 	all_role = frappe.get_doc('Role','All')