If REST client gets an invalid response, regenerate the cached request
headers and try one more time
Signed-off-by: James Ren <jamesren@google.com>
git-svn-id: http://test.kernel.org/svn/autotest/trunk@4392 592f7852-d20e-0410-864c-8624ca9c26a4
diff --git a/frontend/shared/rest_client.py b/frontend/shared/rest_client.py
index b52f213..b5b83af 100644
--- a/frontend/shared/rest_client.py
+++ b/frontend/shared/rest_client.py
@@ -2,6 +2,7 @@
import httplib2
from django.utils import simplejson
from autotest_lib.frontend.afe import rpc_client_lib
+from autotest_lib.client.common_lib import utils
_http = httplib2.Http()
@@ -20,6 +21,16 @@
return headers
+def _clear_request_headers(uri):
+ server = urlparse.urlparse(uri)[0:2]
+ if server in _request_headers:
+ del _request_headers[server]
+
+
+def _site_verify_response_default(headers, response_body):
+ return headers['status'] != '401'
+
+
class RestClientError(Exception):
pass
@@ -116,9 +127,21 @@
logging.debug('%s %s', method, full_uri)
if entity_body:
logging.debug(entity_body)
+
+ site_verify = utils.import_site_function(
+ __file__, 'autotest_lib.frontend.shared.site_rest_client',
+ 'site_verify_response', _site_verify_response_default)
headers, response_body = _http.request(
full_uri, method, body=entity_body,
headers=_get_request_headers(uri))
+ if not site_verify(headers, response_body):
+ logging.debug('Response verification failed, clearing headers and '
+ 'trying again:\n%s', response_body)
+ _clear_request_headers(uri)
+ headers, response_body = _http.request(
+ full_uri, method, body=entity_body,
+ headers=_get_request_headers(uri))
+
logging.debug('Response: %s', headers['status'])
return Response(headers, response_body)