diff --git a/patroni/scripts/aws.py b/patroni/scripts/aws.py
index a0622f15..bdd1bb45 100755
--- a/patroni/scripts/aws.py
+++ b/patroni/scripts/aws.py
@@ -65,8 +65,11 @@ class AWSConnection:
         return self._tag_ebs(new_role) and ret
 
 
-if __name__ == '__main__':
+def main():
     if len(sys.argv) == 4 and sys.argv[1] in ('on_start', 'on_stop', 'on_role_change'):
         AWSConnection(cluster_name=sys.argv[3]).on_role_change(sys.argv[2])
     else:
         sys.exit("Usage: {0} action role name".format(sys.argv[0]))
+
+if __name__ == '__main__':
+    main()
diff --git a/patroni/scripts/wale_restore.py b/patroni/scripts/wale_restore.py
index 368a596c..f54c37a3 100755
--- a/patroni/scripts/wale_restore.py
+++ b/patroni/scripts/wale_restore.py
@@ -139,7 +139,7 @@ class WALERestore(object):
         return ret
 
 
-if __name__ == '__main__':
+def main():
     parser = argparse.ArgumentParser(description='Script to image replicas using WAL-E')
     parser.add_argument('--scope', required=True)
     parser.add_argument('--role', required=False)
@@ -162,3 +162,6 @@ if __name__ == '__main__':
             break
 
     sys.exit(ret)
+
+if __name__ == '__main__':
+    main()