From 5ba34572d9067f69ae4769b1a80be1127d9bd4b3 Mon Sep 17 00:00:00 2001
From: Jeroen Hoekx <jeroen.hoekx@hamok.be>
Date: Fri, 24 Aug 2012 19:08:45 +0200
Subject: [PATCH] Add a state parameter to the wait_for module.

This takes started, stopped and restarted.

Started returns when connecting is possible.
Stopped when connecting is not possible.
Restarted first waits for connecting to be impossible and returns when it is
possible again.
---
 library/wait_for | 44 ++++++++++++++++++++++++++++++++------------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/library/wait_for b/library/wait_for
index 201063d6a0..0c6045e0a4 100644
--- a/library/wait_for
+++ b/library/wait_for
@@ -30,6 +30,7 @@ def main():
             name=dict(required=True),
             timeout=dict(default=300),
             port=dict(default=22),
+            state=dict(default='started', choices=['started', 'stopped', 'restarted']),
         ),
     )
 
@@ -38,21 +39,40 @@ def main():
     host = params['name']
     timeout = int(params['timeout'])
     port = int(params['port'])
+    state = params['state']
 
-    end = datetime.datetime.now() + datetime.timedelta(seconds=timeout)
+    if state in [ 'stopped', 'restarted']:
+        ### first wait for the host to go down
+        end = datetime.datetime.now() + datetime.timedelta(seconds=timeout)
 
-    while datetime.datetime.now() < end:
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        try:
-            s.connect( (host, port) )
-            s.close()
-            break
-        except:
-            time.sleep(1)
-    else:
-        module.fail_json(msg="Timeout when waiting for %s"%(host))
+        while datetime.datetime.now() < end:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.settimeout(5)
+            try:
+                s.connect( (host, port) )
+                s.close()
+                time.sleep(1)
+            except:
+                break
+        else:
+            module.fail_json(msg="Timeout when waiting for %s to stop."%(host))
 
-    module.exit_json(msg="%s responds on %s"%(host, port))
+    if state in [ 'started', 'restarted' ]:
+        ### wait for the host to come up
+        end = datetime.datetime.now() + datetime.timedelta(seconds=timeout)
+
+        while datetime.datetime.now() < end:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            try:
+                s.connect( (host, port) )
+                s.close()
+                break
+            except:
+                time.sleep(1)
+        else:
+            module.fail_json(msg="Timeout when waiting for %s"%(host))
+
+    module.exit_json(msg="State of %s on %s is %s."%(host, port, state))
 
 # this is magic, see lib/ansible/module_common.py
 #<<INCLUDE_ANSIBLE_MODULE_COMMON>>