Merge lp:~allenap/postgresfixture/run-with-command into lp:~lazr-developers/postgresfixture/trunk

[Gavin Panella]

Recast actions as subparsers, and enable run to take a command to execute.

Recast actions as subparsers, and enable run to take a command to execute.

[Graham Binns (gmb)]

review: Approve (code)

=== modified file 'postgresfixture/main.py'
--- postgresfixture/main.py	2012-05-21 10:53:51 +0000
+++ postgresfixture/main.py	2012-05-22 15:17:19 +0000
@@ -62,17 +62,9 @@
     return print(*args, **kwargs)
 
 
-def get_database_name(default="data"):
-    """Return the desired database name, used by some commands.
-
-    Obtained from the ``PGDATABASE`` environment variable.
-    """
-    return environ.get("PGDATABASE", default)
-
-
-def action_destroy(cluster):
-    """Destroy a cluster."""
-    action_stop(cluster)
+def action_destroy(cluster, arguments):
+    """Destroy the cluster."""
+    action_stop(cluster, arguments)
     cluster.destroy()
     if cluster.exists:
         if cluster.lock.locked:
@@ -84,41 +76,34 @@
         raise SystemExit(2)
 
 
-def action_run(cluster):
-    """Create and run a cluster.
-
-    If specified in the ``PGDATABASE`` environment variable, a database will
-    also be created within the cluster.
-    """
-    database_name = get_database_name(default=None)
+def action_run(cluster, arguments):
+    """Create and run the cluster."""
+    database_name = arguments.dbname
+    command = arguments.command
     with cluster:
         if database_name is not None:
             cluster.createdb(database_name)
-        while cluster.running:
-            sleep(5.0)
-
-
-def action_shell(cluster):
-    """Spawn a ``psql`` shell for a database in the cluster.
-
-    The database name can be specified in the ``PGDATABASE`` environment
-    variable.
-    """
-    database_name = get_database_name()
+        if command is None or len(command) == 0:
+            while cluster.running:
+                sleep(5.0)
+        else:
+            cluster.execute(*command)
+
+
+def action_shell(cluster, arguments):
+    """Spawn a `psql` shell for a database in the cluster."""
+    database_name = arguments.dbname
     with cluster:
         cluster.createdb(database_name)
         cluster.shell(database_name)
 
 
-def action_status(cluster):
+def action_status(cluster, arguments):
     """Display a message about the state of the cluster.
 
-    The return code is also set:
-
-    - 0: cluster is running.
-    - 1: cluster exists, but is not running.
-    - 2: cluster does not exist.
-
+    The return code is also set: 0 indicates that the cluster is running; 1
+    indicates that it exists, but is not running; 2 indicates that it does not
+    exist.
     """
     if cluster.exists:
         if cluster.running:
@@ -132,8 +117,8 @@
         raise SystemExit(2)
 
 
-def action_stop(cluster):
-    """Stop a cluster."""
+def action_stop(cluster, arguments):
+    """Stop the cluster."""
     cluster.stop()
     if cluster.running:
         if cluster.lock.locked:
@@ -145,22 +130,10 @@
         raise SystemExit(2)
 
 
-actions = {
-    "destroy": action_destroy,
-    "run": action_run,
-    "shell": action_shell,
-    "status": action_status,
-    "stop": action_stop,
-    }
-
-
 argument_parser = argparse.ArgumentParser(description=__doc__)
 argument_parser.add_argument(
-    "action", choices=sorted(actions), nargs="?", default="shell",
-    help="the action to perform (default: %(default)s)")
-argument_parser.add_argument(
-    "-D", "--datadir", dest="datadir", action="store_true",
-    default="db", help=(
+    "-D", "--datadir", dest="datadir", action="store",
+    metavar="PGDATA", default="db", help=(
         "the directory in which to place, or find, the cluster "
         "(default: %(default)s)"))
 argument_parser.add_argument(
@@ -169,16 +142,57 @@
         "preserve the cluster and its databases when exiting, "
         "even if it was necessary to create and start it "
         "(default: %(default)s)"))
+argument_subparsers = argument_parser.add_subparsers(
+    title="actions")
+
+
+def add_action(name, handler, *args, **kwargs):
+    """Configure a subparser for the given name and function."""
+    parser = argument_subparsers.add_parser(
+        name, *args, help=handler.__doc__, **kwargs)
+    parser.set_defaults(handler=handler)
+    return parser
+
+
+def get_action(name):
+    """Retrieve the named subparser."""
+    return argument_subparsers.choices[name]
+
+
+# Register actions.
+add_action("destroy", action_destroy)
+add_action("run", action_run)
+add_action("shell", action_shell)
+add_action("status", action_status)
+add_action("stop", action_stop)
+
+
+# Customise argument lists for individual actions.
+get_action("run").add_argument(
+    "-d", "--dbname", dest="dbname", action="store", metavar="PGDATABASE",
+    default=environ.get("PGDATABASE", None), help=(
+        "if specified, the database to create. The default is taken from "
+        "the PGDATABASE environment variable (current default: "
+        "%(default)s)."))
+get_action("run").add_argument(
+    "command", nargs="*", default=None, help=(
+        "the command to execute (default: %(default)s)"))
+get_action("shell").add_argument(
+    "-d", "--dbname", dest="dbname", action="store", metavar="PGDATABASE",
+    default=environ.get("PGDATABASE", "data"), help=(
+        "the database to create and connect to. The default is taken from "
+        "the PGDATABASE environment variable, otherwise 'data' (current "
+        "default: %(default)s)."))
 
 
 def main(args=None):
     args = argument_parser.parse_args(args)
     try:
         setup()
-        action = actions[args.action]
         cluster = ClusterFixture(
-            datadir=args.datadir, preserve=args.preserve)
-        action(cluster)
+            datadir=args.datadir,
+            preserve=args.preserve)
+        args.handler(cluster, args)
     except CalledProcessError, error:
         raise SystemExit(error.returncode)
     except KeyboardInterrupt:

=== modified file 'postgresfixture/tests/test_main.py'
--- postgresfixture/tests/test_main.py	2012-05-14 19:19:23 +0000
+++ postgresfixture/tests/test_main.py	2012-05-22 15:17:19 +0000
@@ -32,13 +32,17 @@
     def get_random_database_name(self):
         return "db%s" % b16encode(urandom(8)).lower()
 
+    def parse_args(self, *args):
+        try:
+            return main.argument_parser.parse_args(args)
+        except SystemExit, error:
+            self.fail("parse_args%r failed with %r" % (args, error))
+
     def test_run(self):
         cluster = ClusterFixture(self.make_dir())
         self.addCleanup(cluster.stop)
 
         database_name = self.get_random_database_name()
-        self.useFixture(
-            EnvironmentVariableFixture("PGDATABASE", database_name))
 
         # Instead of sleeping, check the cluster is running, then break out.
         def sleep_patch(time):
@@ -47,7 +51,9 @@
             raise self.Finished
 
         self.patch(main, "sleep", sleep_patch)
-        self.assertRaises(self.Finished, main.action_run, cluster)
+        self.assertRaises(
+            self.Finished, main.action_run, cluster,
+            self.parse_args("run", "--dbname", database_name))
 
     def test_run_without_database(self):
         # A database is not created if it's not specified in the PGDATABASE
@@ -68,22 +74,24 @@
             raise self.Finished
 
         self.patch(main, "sleep", sleep_patch)
-        self.assertRaises(self.Finished, main.action_run, cluster)
+        self.assertRaises(
+            self.Finished, main.action_run, cluster,
+            self.parse_args("run"))
 
     def test_shell(self):
         cluster = ClusterFixture(self.make_dir())
         self.addCleanup(cluster.stop)
 
         database_name = self.get_random_database_name()
-        self.useFixture(
-            EnvironmentVariableFixture("PGDATABASE", database_name))
 
         def shell_patch(database):
             self.assertEqual(database_name, database)
             raise self.Finished
 
         self.patch(cluster, "shell", shell_patch)
-        self.assertRaises(self.Finished, main.action_shell, cluster)
+        self.assertRaises(
+            self.Finished, main.action_shell, cluster,
+            self.parse_args("shell", "--dbname", database_name))
 
     def test_status_running(self):
         cluster = ClusterFixture(self.make_dir())
@@ -91,7 +99,8 @@
         cluster.start()
         self.patch(sys, "stdout", StringIO())
         code = self.assertRaises(
-            SystemExit, main.action_status, cluster).code
+            SystemExit, main.action_status, cluster,
+            self.parse_args("status")).code
         self.assertEqual(0, code)
         self.assertEqual(
             "%s: running\n" % cluster.datadir,
@@ -102,7 +111,8 @@
         cluster.create()
         self.patch(sys, "stdout", StringIO())
         code = self.assertRaises(
-            SystemExit, main.action_status, cluster).code
+            SystemExit, main.action_status, cluster,
+            self.parse_args("status")).code
         self.assertEqual(1, code)
         self.assertEqual(
             "%s: not running\n" % cluster.datadir,
@@ -112,7 +122,8 @@
         cluster = ClusterFixture(self.make_dir())
         self.patch(sys, "stdout", StringIO())
         code = self.assertRaises(
-            SystemExit, main.action_status, cluster).code
+            SystemExit, main.action_status, cluster,
+            self.parse_args("status")).code
         self.assertEqual(2, code)
         self.assertEqual(
             "%s: not created\n" % cluster.datadir,
@@ -122,7 +133,7 @@
         cluster = ClusterFixture(self.make_dir())
         self.addCleanup(cluster.stop)
         cluster.start()
-        main.action_stop(cluster)
+        main.action_stop(cluster, self.parse_args("stop"))
         self.assertFalse(cluster.running)
         self.assertTrue(cluster.exists)
 
@@ -134,7 +145,8 @@
         cluster.lock.acquire()
         self.patch(sys, "stderr", StringIO())
         error = self.assertRaises(
-            SystemExit, main.action_stop, cluster)
+            SystemExit, main.action_stop, cluster,
+            self.parse_args("stop"))
         self.assertEqual(2, error.code)
         self.assertThat(
             sys.stderr.getvalue(), StartsWith(
@@ -145,7 +157,7 @@
         cluster = ClusterFixture(self.make_dir())
         self.addCleanup(cluster.stop)
         cluster.start()
-        main.action_destroy(cluster)
+        main.action_destroy(cluster, self.parse_args("destroy"))
         self.assertFalse(cluster.running)
         self.assertFalse(cluster.exists)
 
@@ -155,7 +167,8 @@
         cluster.lock.acquire()
         self.patch(sys, "stderr", StringIO())
         error = self.assertRaises(
-            SystemExit, main.action_destroy, cluster)
+            SystemExit, main.action_destroy, cluster,
+            self.parse_args("destroy"))
         self.assertEqual(2, error.code)
         self.assertThat(
             sys.stderr.getvalue(), StartsWith(