check_url_available: new allow-redirects option

slapos/promise/plugin/check_url_available.py

@@ -11,6 +11,9 @@ Some notable parameters:
     (default 200) The expected response HTTP code.
   ignore-code:
     (default 0) If set to 1, ignore the response HTTP code.
+  allow-redirects:
+    (default 1) If set to 1, follow Location header on HTTP redirect status code.
+    If set to 0, does not follow and use the redirect response.
   username, password:
     If supplied, enables basic HTTP authentication.
 """
@@ -72,7 +75,7 @@ class RunPromise(GenericPromise):
       request_type = "non-authenticated"
 
     request_options = {
-      'allow_redirects': True,
+      'allow_redirects': bool(int(self.getConfig('allow-redirects', 1))),
       'timeout': int(self.getConfig('timeout', default_timeout)),
       'verify': verify,
       'cert': cert,

slapos/test/promise/plugin/test_check_url_available.py

@@ -191,6 +191,7 @@ class TestHandler(BaseHTTPServer.BaseHTTPRequestHandler):
 
 
 class CheckUrlAvailableMixin(TestPromisePluginMixin):
+  RequestHandler = TestHandler
   @classmethod
   def setUpClass(cls):
     cls.another_server_ca = CertificateAuthority(u"Another Server Root CA")
@@ -221,7 +222,7 @@ class CheckUrlAvailableMixin(TestPromisePluginMixin):
     def server():
       server = BaseHTTPServer.HTTPServer(
         (SLAPOS_TEST_IPV4, SLAPOS_TEST_IPV4_PORT),
-        TestHandler)
+        cls.RequestHandler)
       server.socket = ssl.wrap_socket(
         server.socket,
         certfile=cls.test_server_certificate_file.name,
@@ -630,5 +631,56 @@ class TestCheckUrlAvailableTimeout(CheckUrlAvailableMixin):
     )
 
 
+class TestCheckUrlAvailableRedirect(CheckUrlAvailableMixin):
+  class RequestHandler(BaseHTTPServer.BaseHTTPRequestHandler):
+    def do_GET(self):
+      if self.path == '/':
+        self.send_response(302)
+        self.send_header('Location', '/redirected')
+        self.end_headers()
+        self.wfile.write(b'see /redirected')
+      elif self.path == '/redirected':
+        self.send_response(200)
+        self.end_headers()
+        self.wfile.write(b'OK')
+      else:
+        self.send_response(400)
+        self.end_headers()
+        self.wfile.write(b'Unexepected path: ' + self.path.encode())
+
+  def test_check_redirected_follow_redirect(self):
+    url = HTTPS_ENDPOINT
+    content = self.make_content({
+      'url': url,
+      'http-code': '200',
+    })
+    self.writePromise(self.promise_name, content)
+    self.configureLauncher()
+    self.launcher.run()
+    result = self.getPromiseResult(self.promise_name)
+    self.assertEqual(result['result']['failed'], False)
+    self.assertEqual(
+      result['result']['message'],
+      self.success_template % (url, 200)
+    )
+
+  def test_check_redirected_not_follow_redirect(self):
+    url = HTTPS_ENDPOINT
+    content = self.make_content({
+      'url': url,
+      'allow-redirects': '0',
+      'http-code': '302',
+    })
+    self.writePromise(self.promise_name, content)
+    self.configureLauncher()
+    self.launcher.run()
+    result = self.getPromiseResult(self.promise_name)
+    self.assertEqual(result['result']['failed'], False)
+    self.assertEqual(
+      result['result']['message'],
+      self.success_template % (url, 302)
+    )
+
+
 if __name__ == '__main__':
   unittest.main()