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)