Merge lp:~fwereade/pyjuju/restore-unit-relation-nodes into lp:pyjuju

[William Reade]

When reconstructing unit relation state in UnitLifecycle, ensure presence nodes are created for any relations which are not already departed.

[lp:~fwereade/pyjuju/restore-unit-relation-nodes]

502. By William Reade on 2012-02-03

excise amusing print; expand testing a little

503. By William Reade on 2012-02-03

er, excise the other print

504. By William Reade on 2012-02-08

merge parent

505. By William Reade on 2012-02-21

merge parent

506. By William Reade on 2012-02-21

switch parent to trunk

507. By William Reade on 2012-02-22

merge parent

Unmerged revisions

=== added file 'docs/source/internals/unit-agent-persistence.rst'
--- docs/source/internals/unit-agent-persistence.rst	1970-01-01 00:00:00 +0000
+++ docs/source/internals/unit-agent-persistence.rst	2012-02-02 16:42:42 +0000
@@ -0,0 +1,138 @@
+Notes on unit agent persistence
+===============================
+
+Introduction
+------------
+
+This was first written to explain the extensive changes made in the branch
+lp:~fwereade/juju/restart-transitions; that branch has been split out into
+four separate branches, but this discussion should remain a useful guide to
+the changes made to the unit and relation workflows and lifecycles in the
+course of making the unit agent resumable.
+
+
+Glossary
+--------
+
+UA = UnitAgent
+UL = UnitLifecycle
+UWS = UnitWorkflowState
+URL = UnitRelationLifecycle
+RWS = RelationWorkflowState
+URS = UnitRelationState
+SRS = ServiceRelationState
+
+
+Technical discussion
+--------------------
+
+Probably the most fundamental change is the addition of a "synchronize" method
+to both UWS and RWS. Calling "synchronize" should generally be *all* you need
+to do to put the workflow and associated components into "the right state"; ie
+ZK state will be restored, the appropriate lifecycle will be started (or not),
+and any initial transitons will automatically be fired ("start" for RWS;
+"install", "start" for UWS).
+
+The synchronize method keeps responsibility for the lifecycle's state purely in
+the hands of the workflow; once a workflow is synced, the *only* necessary
+interactions with it should be in response to changes in ZK.
+
+The disadvantage is that lifecycle "start" and "stop" methods have become a
+touch overloaded:
+
+* UL.stop(): now takes "stop_relations" in addition to "fire_hooks", in which
+    "stop_relations" being True causes the orginal behaviour (transition "up"
+    RWSs to "down", as when transitioning the UWS to a "stopped" or "error"
+    state), but False simply causes them to stop watching for changes (in
+    preparation for an orderly shutdown, for example).
+
+* UL.start(): now takes "start_relations" in addition to "fire_hooks", in which
+    the "start_relations" flag being True causes the original behaviour
+    (automatically transition "down" RWSs to "up", as when restarting/resolving
+    the UWS), while False causes the RWSs only to be synced.
+
+* URL.start(): now takes "scheduler" in addition to "watches", allowing the
+    watching and the contained HookScheduler to be controlled separately
+    (allowing us to actually perform the RWS synchronise correctly).
+
+* URL.stop(): still just takes "watches", because there wasn't a scenario in
+    which I wanted to stop the watches but not the HookScheduler.
+
+I still think it's a win, though: and I don't think that turning them into
+separate methods is the right way to go; "start" and "stop" remain perfectly
+decent and appropriate names for what they do.
+
+Now this has been done, we can always launch directly into whatever state we
+shut down in, and that's great, because sudden process death doesn't hurt us
+any more [0] [1]. Except... when we're upgrading a charm. It emerges that the
+charm upgrade state transition only covers the process of firing the hook, and
+not the process of actually upgrading the charm.
+
+In short, we had a mechanism, completely outside the workflow's purview, for
+potentially *brutal* modifications of state (both in terms of the charm itself,
+on disk, and also in that the hook executor should remain stopped forever while
+in "charm_upgrade_error" state); and this rather scuppered the "restart in the
+same state" goal. The obvious thing to do was to move the charm upgrade
+operation into the "charm_upgrade" transition, so we had a *chance* of being
+able to start in the correct state.
+
+UL.upgrade_charm, called by UWS, does itself have subtleties, but it should be
+reasonably clear when examined in context; the most important point is that it
+will call back at the start and end of the risky period, and that the UWS's
+handler for this callback sets a flag in "started"'s state_vars for the
+duration of the upgrade. If that flag is set when we subsequently start up
+again and synchronize the UWS, then we know to immediately force the
+charm_upgrade_error state and work from there.
+
+[0] Well, it does, because we need to persist more than just the (already-
+persisted) workflow state. This branch includes RWS persistence in the UL, as
+requested in this branch's first pre-review (back in the day...), but does not
+include HookScheduler persistence in the URLs, so it remains possible for
+relation hooks which have been queued, but not yet executed, to be lost if the
+process executes before the queue empties. That will be coming in another
+branch (resolve-unit-relation-diffs).
+
+[1] This seems like a good time to mention the UL's relation-broken handling
+for relations that went away while the process was stopped: every time
+._relations is changed, it writes out enough state to recreate a Frankenstein's
+URS object, which it can then use on load to reconstruct the necessary URL and
+hence RWS.
+
+We don't strictly need to *reconstruct* it in every case -- we can just use
+SRS.get_unit_state if the relation still exists -- but given that sometimes we
+do, it seemed senseless to have two code paths for the same operations. Of the
+RWSs we reconstruct, those with existing SRSs will be synchronized (because we
+know it's safe to do so), and the remainder will be stored untouched (because
+we know that _process_service_changes will fire the "depart" transition for us
+before doing anything else... and the "relation-broken" hook will be executed
+in a DepartedRelationHookContext, which is rather restricted, and so shouldn't
+cause the Frankenstein's URS to hit state we can't be sure exists).
+
+
+Appendix: a rough history of changes to restart-transitions
+-----------------------------------------------------------
+
+* Add UWS transitions from "stopped" to "started", so that process restarts can
+    be made to restart UWSs.
+* Upon review, add RWS persistence to UL, to ensure we can't miss
+    relation-broken hooks; as part of this, as discussed, add
+    DepartedRelationHookContext in which to execute them.
+* Upon discussion, discover that original UWS "started" -> "stopped" behaviour
+    on process shutdown is not actually the desired behaviour (and that the
+    associated RWS "up" -> "down" shouldn't happen either.
+* Make changes to UL.start/stop, and add UWS/RWS.synchronize, to allow us to
+    shut down workflows cleanly without transitions and bring them up again in
+    the same state.
+* Discover that we don't have any other reason to transition UWS to "stopped";
+    to actually fire stop hooks at the right time, we need a more sophisticated
+    system (possibly involving the machine agent telling the unit agent to shut
+    itself down). Remove the newly-added "restart" transitions, because they're
+    meaningless now; ponder what good it does us to have a "stopped" state that
+    we never actually enter; chicken out of actually removing it.
+* Realise that charm upgrades do an end-run around the whole UWS mechanism, and
+    resolve to intergate them so I can actually detect upgrades left incomplete
+    due to process death.
+* Move charm upgrade operation from agent into UL; come to appreciate the
+    subtleties of the charm upgrade process; make necessary tweaks to
+    UL.upgrade_charm, and UWS, to allow for synchronization of incomplete
+    upgrades.

=== modified file 'examples/oneiric/mysql/hooks/install'
--- examples/oneiric/mysql/hooks/install	2011-09-15 18:56:08 +0000
+++ examples/oneiric/mysql/hooks/install	2012-02-02 16:42:42 +0000
@@ -21,5 +21,5 @@
 juju-log "Editing my.cnf to allow listening on all interfaces"
 sed --in-place=old 's/127\.0\.0\.1/0.0.0.0/' /etc/mysql/my.cnf
 
-juju-log "Restarting mysql service"
-service mysql restart
+juju-log "Stopping mysql service"
+service mysql stop

=== modified file 'examples/oneiric/mysql/hooks/start'
--- examples/oneiric/mysql/hooks/start	2011-02-03 01:23:43 +0000
+++ examples/oneiric/mysql/hooks/start	2012-02-02 16:42:42 +0000
@@ -1,1 +1,3 @@
-#!/bin/bash
\ No newline at end of file
+#!/bin/bash
+juju-log "Starting mysql service"
+service mysql start || service mysql restart

=== modified file 'examples/oneiric/mysql/hooks/stop'
--- examples/oneiric/mysql/hooks/stop	2011-09-15 18:56:08 +0000
+++ examples/oneiric/mysql/hooks/stop	2012-02-02 16:42:42 +0000
@@ -1,3 +1,3 @@
 #!/bin/bash
 juju-log "Stopping mysql service"
-/etc/init.d/mysql stop
+service mysql stop || true

=== modified file 'juju/agents/base.py'
--- juju/agents/base.py	2011-09-22 13:23:00 +0000
+++ juju/agents/base.py	2012-02-02 16:42:42 +0000
@@ -1,7 +1,9 @@
 import argparse
 import os
+import logging
+import stat
 import sys
-import logging
+import yaml
 
 import zookeeper
 
@@ -18,6 +20,23 @@
 from juju.state.environment import GlobalSettingsStateManager
 
 
+def load_client_id(path):
+    try:
+        with open(path) as f:
+            return yaml.load(f.read())
+    except IOError:
+        return None
+
+
+def save_client_id(path, client_id):
+    parent = os.path.dirname(path)
+    if not os.path.exists(parent):
+        os.makedirs(parent)
+    with open(path, "w") as f:
+        f.write(yaml.dump(client_id))
+    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
+
+
 class TwistedOptionNamespace(object):
     """
     An argparse namespace implementation that is compatible with twisted
@@ -153,13 +172,40 @@
                 "Invalid juju-directory %r, does not exist." % (
                     options.get("juju_directory")))
 
+        if options["session_file"] is None:
+            raise JujuError("No session file specified")
+
         self.config = options
 
     @inlineCallbacks
+    def _kill_existing_session(self):
+        try:
+            # We might have died suddenly, in which case the session may
+            # still be alive. If this is the case, shoot it in the head, so
+            # it doesn't interfere with our attempts to recreate our state.
+            # (We need to be able to recreate our state *anyway*, and it's
+            # much simpler to force ourselves to recreate it every time than
+            # it is to mess around partially recreating partial state.)
+            client_id = load_client_id(self.config["session_file"])
+            if client_id is None:
+                return
+            temp_client = yield ZookeeperClient().connect(
+                self.config["zookeeper_servers"], client_id=client_id)
+            yield temp_client.close()
+        except zookeeper.ZooKeeperException:
+            # We don't really care what went wrong; just that we're not able
+            # to connect using the old session, and therefore we should be ok
+            # to start a fresh one without transient state hanging around.
+            pass
+
+    @inlineCallbacks
     def connect(self):
         """Return an authenticated connection to the juju zookeeper."""
-        hosts = self.config["zookeeper_servers"]
-        self.client = yield ZookeeperClient().connect(hosts)
+        yield self._kill_existing_session()
+        self.client = yield ZookeeperClient().connect(
+            self.config["zookeeper_servers"])
+        save_client_id(
+            self.config["session_file"], self.client.client_id)
 
         principals = self.config.get("principals", ())
         for principal in principals:
@@ -200,6 +246,9 @@
         finally:
             if self.client and self.client.connected:
                 self.client.close()
+            session_file = self.config["session_file"]
+            if os.path.exists(session_file):
+                os.unlink(session_file)
 
     def set_watch_enabled(self, flag):
         """Set boolean flag for whether this agent should watching zookeeper.
@@ -285,3 +334,7 @@
     parser.add_argument(
         "--juju-directory", default=juju_home, type=os.path.abspath,
         help="juju working directory ($JUJU_HOME)")
+
+    parser.add_argument(
+        "--session-file", default=None, type=os.path.abspath,
+        help="like a pidfile, but for the zookeeper session id")

=== modified file 'juju/agents/tests/common.py'
--- juju/agents/tests/common.py	2011-09-15 19:24:47 +0000
+++ juju/agents/tests/common.py	2012-02-02 16:42:42 +0000
@@ -55,6 +55,7 @@
         options = TwistedOptionNamespace()
         options["juju_directory"] = self.juju_directory
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         return succeed(options)
 
     @inlineCallbacks

=== modified file 'juju/agents/tests/test_base.py'
--- juju/agents/tests/test_base.py	2011-09-15 18:50:23 +0000
+++ juju/agents/tests/test_base.py	2012-02-02 16:42:42 +0000
@@ -2,13 +2,14 @@
 import json
 import logging
 import os
+import stat
 import sys
-
+import yaml
 
 from twisted.application.app import AppLogger
 from twisted.application.service import IService, IServiceCollection
 from twisted.internet.defer import (
-    succeed, Deferred, inlineCallbacks, returnValue)
+    fail, succeed, Deferred, inlineCallbacks, returnValue)
 from twisted.python.components import Componentized
 from twisted.python import log
 
@@ -20,6 +21,7 @@
 
 from juju.agents.base import (
     BaseAgent, TwistedOptionNamespace, AgentRunner, AgentLogger)
+from juju.agents.dummy import DummyAgent
 from juju.errors import NoConnection, JujuError
 from juju.lib.zklog import ZookeeperHandler
 
@@ -34,8 +36,8 @@
     @inlineCallbacks
     def setUp(self):
         yield super(BaseAgentTest, self).setUp()
-        self.change_environment(
-            JUJU_HOME=self.makeDir())
+        self.juju_home = self.makeDir()
+        self.change_environment(JUJU_HOME=self.juju_home)
 
     def test_as_app(self):
         """The agent class can be accessed as an application."""
@@ -53,11 +55,10 @@
         # Daemon group
         self.assertEqual(
             parser.get_default("logfile"), "%s.log" % BaseAgent.name)
-        self.assertEqual(
-            parser.get_default("pidfile"), "%s.pid" % BaseAgent.name)
+        self.assertEqual(parser.get_default("pidfile"), "")
 
         self.assertEqual(parser.get_default("loglevel"), "DEBUG")
-        self.assertTrue(parser.get_default("nodaemon"))
+        self.assertFalse(parser.get_default("nodaemon"))
         self.assertEqual(parser.get_default("rundir"), ".")
         self.assertEqual(parser.get_default("chroot"), None)
         self.assertEqual(parser.get_default("umask"), None)
@@ -80,6 +81,8 @@
         # Agent options
         self.assertEqual(parser.get_default("principals"), [])
         self.assertEqual(parser.get_default("zookeeper_servers"), "")
+        self.assertEqual(parser.get_default("juju_directory"), self.juju_home)
+        self.assertEqual(parser.get_default("session_file"), None)
 
     def test_twistd_flags_correspond(self):
         parser = argparse.ArgumentParser()
@@ -87,11 +90,11 @@
         args = [
             "--profile",
             "--savestats",
-            "--daemon"]
+            "--nodaemon"]
 
         options = parser.parse_args(args, namespace=TwistedOptionNamespace())
         self.assertEqual(options.get("savestats"), True)
-        self.assertEqual(options.get("nodaemon"), False)
+        self.assertEqual(options.get("nodaemon"), True)
         self.assertEqual(options.get("profile"), True)
 
     def test_agent_logger(self):
@@ -100,7 +103,8 @@
         log_file_path = self.makeFile()
 
         options = parser.parse_args(
-            ["--logfile", log_file_path], namespace=TwistedOptionNamespace())
+            ["--logfile", log_file_path, "--session-file", self.makeFile()],
+            namespace=TwistedOptionNamespace())
 
         def match_observer(observer):
             return isinstance(observer.im_self, log.PythonLoggingObserver)
@@ -181,8 +185,9 @@
         This will create an agent instance,  parse the cli args, passes them to
         the agent, and starts the agent runner.
         """
-        self.change_args("es-agent", "--zookeeper-servers",
-                         get_test_zookeeper_address())
+        self.change_args(
+            "es-agent", "--zookeeper-servers", get_test_zookeeper_address(),
+            "--session-file", self.makeFile())
         runner = self.mocker.patch(AgentRunner)
         runner.run()
         mock_agent = self.mocker.patch(BaseAgent)
@@ -219,11 +224,10 @@
 
         started.addCallback(validate_started)
 
-        pid_file = self.makeFile()
         self.change_args(
-            "es-agent",
+            "es-agent", "--nodaemon",
             "--zookeeper-servers", get_test_zookeeper_address(),
-            "--pidfile", pid_file)
+            "--session-file", self.makeFile())
         runner = self.mocker.patch(AgentRunner)
         logger = self.mocker.patch(AppLogger)
         logger.start(MATCH_APP)
@@ -233,6 +237,7 @@
         DummyAgent.run()
         return started
 
+    @inlineCallbacks
     def test_stop_service_stub_closes_agent(self):
         """The base class agent, stopService will the stop method.
 
@@ -241,6 +246,7 @@
         """
         mock_agent = self.mocker.patch(BaseAgent)
         mock_client = self.mocker.mock(ZookeeperClient)
+        session_file = self.makeFile()
 
         # connection is closed after agent.stop invoked.
         with self.mocker.order():
@@ -262,11 +268,17 @@
             self.mocker.result(mock_client)
             mock_client.close()
 
+            # delete session file
+            mock_agent.config
+            self.mocker.result({"session_file": session_file})
+
         self.mocker.replay()
 
         agent = BaseAgent()
-        return agent.stopService()
+        yield agent.stopService()
+        self.assertFalse(os.path.exists(session_file))
 
+    @inlineCallbacks
     def test_stop_service_stub_ignores_disconnected_agent(self):
         """The base class agent, stopService will the stop method.
 
@@ -274,6 +286,7 @@
         """
         mock_agent = self.mocker.patch(BaseAgent)
         mock_client = self.mocker.mock(ZookeeperClient)
+        session_file = self.makeFile()
 
         # connection is closed after agent.stop invoked.
         with self.mocker.order():
@@ -289,10 +302,14 @@
             mock_client.connected
             self.mocker.result(False)
 
+            mock_agent.config
+            self.mocker.result({"session_file": session_file})
+
         self.mocker.replay()
 
         agent = BaseAgent()
-        return agent.stopService()
+        yield agent.stopService()
+        self.assertFalse(os.path.exists(session_file))
 
     def test_run_base_raises_error(self):
         """The base class agent, raises a notimplemented error when started."""
@@ -300,12 +317,15 @@
         client.connect(get_test_zookeeper_address())
         client_mock = self.mocker.mock()
         self.mocker.result(succeed(client_mock))
+        client_mock.client_id
+        self.mocker.result((123, "abc"))
         self.mocker.replay()
 
         agent = BaseAgent()
         agent.configure({
             "zookeeper_servers": get_test_zookeeper_address(),
-            "juju_directory": self.makeDir()})
+            "juju_directory": self.makeDir(),
+            "session_file": self.makeFile()})
         d = agent.startService()
         self.failUnlessFailure(d, NotImplementedError)
         return d
@@ -316,35 +336,43 @@
         client = self.mocker.patch(ZookeeperClient)
         client.connect("x2.example.com")
         self.mocker.result(succeed(mock_client))
+        mock_client.client_id
+        self.mocker.result((123, "abc"))
         self.mocker.replay()
 
         agent = BaseAgent()
         agent.configure({"zookeeper_servers": "x2.example.com",
-                         "juju_directory": self.makeDir()})
+                         "juju_directory": self.makeDir(),
+                         "session_file": self.makeFile()})
         result = agent.connect()
         self.assertEqual(result.result, mock_client)
         self.assertEqual(agent.client, mock_client)
 
-    def test_non_existant_directory(self):
+    def test_nonexistent_directory(self):
         """If the juju directory does not exist an error should be raised.
         """
         juju_directory = self.makeDir()
         os.rmdir(juju_directory)
         data = {"zookeeper_servers": get_test_zookeeper_address(),
-                "juju_directory": juju_directory}
+                "juju_directory": juju_directory,
+                "session_file": self.makeFile()}
+        self.assertRaises(JujuError, BaseAgent().configure, data)
 
-        agent = BaseAgent()
-        self.assertRaises(
-            JujuError,
-            agent.configure,
-            data)
+    def test_bad_session_file(self):
+        """If the session file cannot be created an error should be raised.
+        """
+        data = {"zookeeper_servers": get_test_zookeeper_address(),
+                "juju_directory": self.makeDir(),
+                "session_file": None}
+        self.assertRaises(JujuError, BaseAgent().configure, data)
 
     def test_directory_cli_option(self):
         """The juju directory can be configured on the cli."""
         juju_directory = self.makeDir()
         self.change_args(
             "es-agent", "--zookeeper-servers", get_test_zookeeper_address(),
-            "--juju-directory", juju_directory)
+            "--juju-directory", juju_directory,
+            "--session-file", self.makeFile())
 
         agent = BaseAgent()
         parser = argparse.ArgumentParser()
@@ -366,7 +394,9 @@
         agent = BaseAgent()
         parser = argparse.ArgumentParser()
         agent.setup_options(parser)
-        options = parser.parse_args(namespace=TwistedOptionNamespace())
+        options = parser.parse_args(
+            ["--session-file", self.makeFile()],
+            namespace=TwistedOptionNamespace())
         agent.configure(options)
         self.assertEqual(
             agent.config["juju_directory"], juju_directory)
@@ -382,6 +412,8 @@
         client = self.mocker.patch(ZookeeperClient)
         client.connect("x1.example.com")
         self.mocker.result(succeed(client))
+        client.client_id
+        self.mocker.result((123, "abc"))
         client.add_auth("digest", "admin:abc")
         client.add_auth("digest", "agent:xyz")
         client.exists("/")
@@ -390,7 +422,105 @@
         agent = BaseAgent()
         parser = argparse.ArgumentParser()
         agent.setup_options(parser)
-        options = parser.parse_args(namespace=TwistedOptionNamespace())
+        options = parser.parse_args(
+            ["--session-file", self.makeFile()],
+            namespace=TwistedOptionNamespace())
+        agent.configure(options)
+        d = agent.startService()
+        self.failUnlessFailure(d, NotImplementedError)
+        return d
+
+    def test_connect_closes_running_session(self):
+        self.change_args("es-agent")
+        self.change_environment(
+            JUJU_HOME=self.makeDir(),
+            JUJU_ZOOKEEPER="x1.example.com")
+
+        session_file = self.makeFile()
+        with open(session_file, "w") as f:
+            f.write(yaml.dump((123, "abc")))
+        mock_client_1 = self.mocker.mock()
+        client = self.mocker.patch(ZookeeperClient)
+        client.connect("x1.example.com", client_id=(123, "abc"))
+        self.mocker.result(succeed(mock_client_1))
+        mock_client_1.close()
+        self.mocker.result(None)
+
+        mock_client_2 = self.mocker.mock()
+        client.connect("x1.example.com")
+        self.mocker.result(succeed(mock_client_2))
+        mock_client_2.client_id
+        self.mocker.result((456, "def"))
+        self.mocker.replay()
+
+        agent = BaseAgent()
+        parser = argparse.ArgumentParser()
+        agent.setup_options(parser)
+        options = parser.parse_args(
+            ["--session-file", session_file],
+            namespace=TwistedOptionNamespace())
+        agent.configure(options)
+        d = agent.startService()
+        self.failUnlessFailure(d, NotImplementedError)
+        return d
+
+    def test_connect_handles_expired_session(self):
+        self.change_args("es-agent")
+        self.change_environment(
+            JUJU_HOME=self.makeDir(),
+            JUJU_ZOOKEEPER="x1.example.com")
+
+        session_file = self.makeFile()
+        with open(session_file, "w") as f:
+            f.write(yaml.dump((123, "abc")))
+        client = self.mocker.patch(ZookeeperClient)
+        client.connect("x1.example.com", client_id=(123, "abc"))
+        self.mocker.result(fail(zookeeper.SessionExpiredException()))
+
+        mock_client = self.mocker.mock()
+        client.connect("x1.example.com")
+        self.mocker.result(succeed(mock_client))
+        mock_client.client_id
+        self.mocker.result((456, "def"))
+        self.mocker.replay()
+
+        agent = BaseAgent()
+        parser = argparse.ArgumentParser()
+        agent.setup_options(parser)
+        options = parser.parse_args(
+            ["--session-file", session_file],
+            namespace=TwistedOptionNamespace())
+        agent.configure(options)
+        d = agent.startService()
+        self.failUnlessFailure(d, NotImplementedError)
+        return d
+
+    def test_connect_handles_nonsense_session(self):
+        self.change_args("es-agent")
+        self.change_environment(
+            JUJU_HOME=self.makeDir(),
+            JUJU_ZOOKEEPER="x1.example.com")
+
+        session_file = self.makeFile()
+        with open(session_file, "w") as f:
+            f.write(yaml.dump("cheesy wotsits"))
+        client = self.mocker.patch(ZookeeperClient)
+        client.connect("x1.example.com", client_id="cheesy wotsits")
+        self.mocker.result(fail(zookeeper.ZooKeeperException()))
+
+        mock_client = self.mocker.mock()
+        client.connect("x1.example.com")
+        self.mocker.result(succeed(mock_client))
+        mock_client.client_id
+        self.mocker.result((456, "def"))
+        self.mocker.replay()
+
+        agent = BaseAgent()
+        parser = argparse.ArgumentParser()
+        agent.setup_options(parser)
+        options = parser.parse_args(
+            ["--session-file", session_file],
+            namespace=TwistedOptionNamespace())
         agent.configure(options)
         d = agent.startService()
         self.failUnlessFailure(d, NotImplementedError)
@@ -408,6 +538,21 @@
         agent.set_watch_enabled(False)
         self.assertFalse(agent.get_watch_enabled())
 
+    @inlineCallbacks
+    def test_session_file_permissions(self):
+        session_file = self.makeFile()
+        agent = DummyAgent()
+        agent.configure({
+            "session_file": session_file,
+            "juju_directory": self.makeDir(),
+            "zookeeper_servers": get_test_zookeeper_address()})
+        yield agent.startService()
+        mode = os.stat(session_file).st_mode
+        mask = stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO
+        self.assertEquals(mode & mask, stat.S_IRUSR | stat.S_IWUSR)
+        yield agent.stopService()
+        self.assertFalse(os.path.exists(session_file))
+
 
 class AgentDebugLogSettingsWatch(AgentTestBase):
 

=== modified file 'juju/agents/tests/test_machine.py'
--- juju/agents/tests/test_machine.py	2011-10-05 13:59:44 +0000
+++ juju/agents/tests/test_machine.py	2012-02-02 16:42:42 +0000
@@ -120,10 +120,9 @@
         # initially setup by get_agent_config in setUp
         self.change_environment(JUJU_MACHINE_ID="")
         self.change_args("es-agent",
-                         "--zookeeper-servers",
-                         get_test_zookeeper_address(),
-                         "--juju-directory",
-                         self.makeDir())
+                         "--zookeeper-servers", get_test_zookeeper_address(),
+                         "--juju-directory", self.makeDir(),
+                         "--session-file", self.makeFile())
         parser = argparse.ArgumentParser()
         self.agent.setup_options(parser)
         options = parser.parse_args(namespace=TwistedOptionNamespace())

=== modified file 'juju/agents/tests/test_unit.py'
--- juju/agents/tests/test_unit.py	2011-12-16 09:23:31 +0000
+++ juju/agents/tests/test_unit.py	2012-02-02 16:42:42 +0000
@@ -3,10 +3,9 @@
 import os
 import yaml
 
-from twisted.internet.defer import (
-    inlineCallbacks, returnValue, fail, Deferred)
+from twisted.internet.defer import inlineCallbacks, returnValue
 
-from juju.agents.unit import UnitAgent, CharmUpgradeOperation
+from juju.agents.unit import UnitAgent
 from juju.agents.base import TwistedOptionNamespace
 from juju.charm import get_charm_from_path
 from juju.charm.url import CharmURL
@@ -74,8 +73,10 @@
                 "stop", "#!/bin/bash\necho stop >> %s" % output_file)
 
         for k in kw.keys():
-            self.write_hook(k.replace("_", "-"),
-                            "#!/bin/bash\necho $0 >> %s" % output_file)
+            hook_name = k.replace("_", "-")
+            self.write_hook(
+                hook_name,
+                "#!/bin/bash\necho %s >> %s" % (hook_name, output_file))
 
         return output_file
 
@@ -183,7 +184,8 @@
         self.change_args(
             "unit-agent",
             "--juju-directory", self.makeDir(),
-            "--zookeeper-servers", get_test_zookeeper_address())
+            "--zookeeper-servers", get_test_zookeeper_address(),
+            "--session-file", self.makeFile())
 
         parser = argparse.ArgumentParser()
         self.agent.setup_options(parser)
@@ -215,6 +217,7 @@
         options = {}
         options["juju_directory"] = self.juju_directory
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         options["unit_name"] = "rabbit-1"
         agent = self.agent_class()
         agent.configure(options)
@@ -568,6 +571,12 @@
         self.makeDir(path=os.path.join(self.juju_directory, "charms"))
 
     @inlineCallbacks
+    def wait_for_log(self, logger_name, message, level=logging.DEBUG):
+        output = self.capture_logging(logger_name, level=level)
+        while message not in output.getvalue():
+            yield self.sleep(0.1)
+
+    @inlineCallbacks
     def mark_charm_upgrade(self):
         # Create a new version of the charm
         repository = self.increment_charm(self.charm)
@@ -592,158 +601,84 @@
         yield self.assertState(self.agent.workflow, "started")
 
     @inlineCallbacks
-    def test_agent_upgrade_watch_continues_on_unexpected_error(self):
-        """The agent watches for unit upgrades and continues if there is an
-        unexpected error."""
-        yield self.mark_charm_upgrade()
-        self.agent.set_watch_enabled(True)
-
-        output = self.capture_logging(
-            "juju.agents.unit", level=logging.DEBUG)
-
-        upgrade_done = Deferred()
-
-        def operation_has_run():
-            upgrade_done.callback(True)
-
-        operation = self.mocker.patch(CharmUpgradeOperation)
-        operation.run()
-
-        self.mocker.call(operation_has_run)
-        self.mocker.result(fail(ValueError("magic mouse")))
-        self.mocker.replay()
-
-        yield self.agent.startService()
-
-        yield upgrade_done
-        self.assertIn("Error while upgrading", output.getvalue())
-        self.assertIn("magic mouse", output.getvalue())
-
-        yield self.agent.workflow.fire_transition("stop")
-
-    @inlineCallbacks
     def test_agent_upgrade(self):
         """The agent can succesfully upgrade its charm."""
-        self.agent.set_watch_enabled(False)
-        yield self.agent.startService()
-
-        yield self.mark_charm_upgrade()
-
+        log_written = self.wait_for_log("juju.agents.unit", "Upgrade complete")
         hook_done = self.wait_on_hook(
             "upgrade-charm", executor=self.agent.executor)
-        self.write_hook("upgrade-charm", "#!/bin/bash\nexit 0")
-        output = self.capture_logging("unit.upgrade", level=logging.DEBUG)
-
-        # Do the upgrade
-        upgrade = CharmUpgradeOperation(self.agent)
-        value = yield upgrade.run()
-
-        # Verify the upgrade.
-        self.assertIdentical(value, True)
-        self.assertIn("Unit upgraded", output.getvalue())
+
+        self.agent.set_watch_enabled(True)
+        yield self.agent.startService()
+        yield self.mark_charm_upgrade()
         yield hook_done
+        yield log_written
 
+        self.assertIdentical(
+            (yield self.states["unit"].get_upgrade_flag()),
+            False)
         new_charm = get_charm_from_path(
             os.path.join(self.agent.unit_directory, "charm"))
-
         self.assertEqual(
             self.charm.get_revision() + 1, new_charm.get_revision())
 
     @inlineCallbacks
+    def test_agent_upgrade_version_current(self):
+        """If the unit is running the latest charm, do nothing."""
+        log_written = self.wait_for_log(
+            "juju.agents.unit",
+            "Upgrade ignored: already running latest charm")
+
+        old_charm_id = yield self.states["unit"].get_charm_id()
+        self.agent.set_watch_enabled(True)
+        yield self.agent.startService()
+        yield self.states["unit"].set_upgrade_flag()
+        yield log_written
+
+        self.assertIdentical(
+            (yield self.states["unit"].get_upgrade_flag()), False)
+        self.assertEquals(
+            (yield self.states["unit"].get_charm_id()), old_charm_id)
+
+
+    @inlineCallbacks
     def test_agent_upgrade_bad_unit_state(self):
-        """The an upgrade fails if the unit is in a bad state."""
-        self.agent.set_watch_enabled(False)
-        yield self.agent.startService()
-
+        """The upgrade fails if the unit is in a bad state."""
         # Upload a new version of the unit's charm
         repository = self.increment_charm(self.charm)
         charm = yield repository.find(CharmURL.parse("local:series/mysql"))
         charm, charm_state = yield self.publish_charm(charm.path)
+        old_charm_id = yield self.states["unit"].get_charm_id()
+
+        log_written = self.wait_for_log(
+            "juju.agents.unit",
+            "Cannot upgrade: unit is in non-started state configure_error. "
+            "Reissue upgrade command to try again.")
+        self.agent.set_watch_enabled(True)
+        yield self.agent.startService()
 
         # Mark the unit for upgrade, with an invalid state.
+        yield self.agent.workflow.fire_transition("error_configure")
         yield self.states["service"].set_charm_id(charm_state.id)
         yield self.states["unit"].set_upgrade_flag()
-        yield self.agent.workflow.set_state("start_error")
-
-        output = self.capture_logging("unit.upgrade", level=logging.DEBUG)
-
-        # Do the upgrade
-        upgrade = CharmUpgradeOperation(self.agent)
-        value = yield upgrade.run()
-
-        # Verify the upgrade.
-        self.assertIdentical(value, False)
-        self.assertIn("Unit not in an upgradeable state: start_error",
-                      output.getvalue())
+        yield log_written
+
         self.assertIdentical(
-            (yield self.states["unit"].get_upgrade_flag()),
-            False)
+            (yield self.states["unit"].get_upgrade_flag()), False)
+        self.assertEquals(
+            (yield self.states["unit"].get_charm_id()), old_charm_id)
 
     @inlineCallbacks
     def test_agent_upgrade_no_flag(self):
-        """An upgrade fails if there is no upgrade flag set."""
-        self.agent.set_watch_enabled(False)
-        yield self.agent.startService()
-        output = self.capture_logging("unit.upgrade", level=logging.DEBUG)
-        upgrade = CharmUpgradeOperation(self.agent)
-        value = yield upgrade.run()
-        self.assertIdentical(value, False)
-        self.assertIn("No upgrade flag set", output.getvalue())
-        yield self.agent.workflow.fire_transition("stop")
-
-    @inlineCallbacks
-    def test_agent_upgrade_version_current(self):
-        """An upgrade fails if the unit is running the latest charm."""
-        self.agent.set_watch_enabled(False)
-        yield self.agent.startService()
-        yield self.states["unit"].set_upgrade_flag()
-        output = self.capture_logging("unit.upgrade", level=logging.DEBUG)
-        upgrade = CharmUpgradeOperation(self.agent)
-        value = yield upgrade.run()
-        self.assertIdentical(value, True)
-        self.assertIn("Unit already running latest charm", output.getvalue())
-        self.assertFalse((yield self.states["unit"].get_upgrade_flag()))
-
-    @inlineCallbacks
-    def test_agent_upgrade_hook_failure(self):
-        """An upgrade fails if the upgrade hook errors."""
-        self.agent.set_watch_enabled(False)
-        yield self.agent.startService()
-
-        # Upload a new version of the unit's charm
-        repository = self.increment_charm(self.charm)
-        charm = yield repository.find(CharmURL.parse("local:series/mysql"))
-        charm, charm_state = yield self.publish_charm(charm.path)
-
-        # Mark the unit for upgrade
-        yield self.states["service"].set_charm_id(charm_state.id)
-        yield self.states["unit"].set_upgrade_flag()
-
-        hook_done = self.wait_on_hook(
-            "upgrade-charm", executor=self.agent.executor)
-        self.write_hook("upgrade-charm", "#!/bin/bash\nexit 1")
-        output = self.capture_logging("unit.upgrade", level=logging.DEBUG)
-
-        # Do the upgrade
-        upgrade = CharmUpgradeOperation(self.agent)
-        value = yield upgrade.run()
-
-        # Verify the failed upgrade.
-        self.assertIdentical(value, False)
-        self.assertIn("Invoking upgrade transition", output.getvalue())
-        self.assertIn("Upgrade failed.", output.getvalue())
-        yield hook_done
-
-        # Verify state
-        workflow_state = yield self.agent.workflow.get_state()
-        self.assertEqual("charm_upgrade_error", workflow_state)
-
-        # Verify new charm is in place
-        new_charm = get_charm_from_path(
-            os.path.join(self.agent.unit_directory, "charm"))
-
-        self.assertEqual(
-            self.charm.get_revision() + 1, new_charm.get_revision())
-
-        # Verify upgrade flag is cleared.
-        self.assertFalse((yield self.states["unit"].get_upgrade_flag()))
+        """An upgrade stops if there is no upgrade flag set."""
+        log_written = self.wait_for_log(
+            "juju.agents.unit", "No upgrade flag set")
+        old_charm_id = yield self.states["unit"].get_charm_id()
+        self.agent.set_watch_enabled(True)
+        yield self.agent.startService()
+        yield log_written
+
+        self.assertIdentical(
+            (yield self.states["unit"].get_upgrade_flag()),
+            False)
+        new_charm_id = yield self.states["unit"].get_charm_id()
+        self.assertEquals(new_charm_id, old_charm_id)

=== modified file 'juju/agents/unit.py'
--- juju/agents/unit.py	2012-01-10 14:14:28 +0000
+++ juju/agents/unit.py	2012-02-02 16:42:42 +0000
@@ -1,7 +1,5 @@
 import os
 import logging
-import shutil
-import tempfile
 
 from twisted.internet.defer import inlineCallbacks, returnValue
 
@@ -14,8 +12,6 @@
 from juju.unit.lifecycle import UnitLifecycle, HOOK_SOCKET_FILE
 from juju.unit.workflow import UnitWorkflowState
 
-from juju.unit.charm import download_charm
-
 from juju.agents.base import BaseAgent
 
 
@@ -66,14 +62,14 @@
     @inlineCallbacks
     def start(self):
         """Start the unit agent process."""
-        self.service_state_manager = ServiceStateManager(self.client)
+        service_state_manager = ServiceStateManager(self.client)
 
         # Retrieve our unit and configure working directories.
         service_name = self.unit_name.split("/")[0]
-        service_state = yield self.service_state_manager.get_service_state(
+        self.service_state = yield service_state_manager.get_service_state(
             service_name)
 
-        self.unit_state = yield service_state.get_unit_state(
+        self.unit_state = yield self.service_state.get_unit_state(
             self.unit_name)
         self.unit_directory = os.path.join(
             self.config["juju_directory"], "units",
@@ -82,10 +78,11 @@
             self.config["juju_directory"], "state")
 
         # Setup the server portion of the cli api exposed to hooks.
+        socket_path = os.path.join(self.unit_directory, HOOK_SOCKET_FILE)
+        if os.path.exists(socket_path):
+            os.unlink(socket_path)
         from twisted.internet import reactor
-        self.api_socket = reactor.listenUNIX(
-            os.path.join(self.unit_directory, HOOK_SOCKET_FILE),
-            self.api_factory)
+        self.api_socket = reactor.listenUNIX(socket_path, self.api_factory)
 
         # Setup the unit state's address
         address = yield get_unit_address(self.client)
@@ -100,9 +97,13 @@
         # Inform the system, we're alive.
         yield self.unit_state.connect_agent()
 
+        # Start paying attention to the debug-log setting
+        if self.get_watch_enabled():
+            yield self.unit_state.watch_hook_debug(self.cb_watch_hook_debug)
+
         self.lifecycle = UnitLifecycle(
-            self.client, self.unit_state, service_state, self.unit_directory,
-            self.state_directory, self.executor)
+            self.client, self.unit_state, self.service_state,
+            self.unit_directory, self.state_directory, self.executor)
 
         self.workflow = UnitWorkflowState(
             self.client, self.unit_state, self.lifecycle, self.state_directory)
@@ -113,7 +114,7 @@
 
         if self.get_watch_enabled():
             yield self.unit_state.watch_resolved(self.cb_watch_resolved)
-            yield service_state.watch_config_state(
+            yield self.service_state.watch_config_state(
                 self.cb_watch_config_changed)
             yield self.unit_state.watch_upgrade_flag(
                 self.cb_watch_upgrade_flag)
@@ -175,13 +176,34 @@
         """Update the unit's charm when requested.
         """
         upgrade_flag = yield self.unit_state.get_upgrade_flag()
-        if upgrade_flag:
-            log.info("Upgrade detected, starting upgrade")
-            upgrade = CharmUpgradeOperation(self)
-            try:
-                yield upgrade.run()
-            except Exception:
-                log.exception("Error while upgrading")
+        if not upgrade_flag:
+            log.info("No upgrade flag set.")
+            return
+
+        log.info("Upgrade detected")
+        # Clear the flag immediately; this means that upgrade requests will
+        # be *ignored* by units which are not "started", and will need to be
+        # reissued when the units are in acceptable states.
+        yield self.unit_state.clear_upgrade_flag()
+
+        new_id = yield self.service_state.get_charm_id()
+        old_id = yield self.unit_state.get_charm_id()
+        if new_id == old_id:
+            log.info("Upgrade ignored: already running latest charm")
+            return
+
+        state = yield self.workflow.get_state()
+        if state != "started":
+            log.warning(
+                "Cannot upgrade: unit is in non-started state %s. Reissue "
+                "upgrade command to try again.", state)
+            return
+
+        log.info("Starting upgrade")
+        if (yield self.workflow.fire_transition("upgrade_charm")):
+            log.info("Upgrade complete")
+        else:
+            log.info("Upgrade failed")
 
     @inlineCallbacks
     def cb_watch_config_changed(self, change):
@@ -198,99 +220,5 @@
         yield self.workflow.fire_transition("reconfigure")
 
 
-class CharmUpgradeOperation(object):
-    """A unit agent charm upgrade operation."""
-
-    def __init__(self, agent):
-        self._agent = agent
-        self._log = logging.getLogger("unit.upgrade")
-        self._charm_directory = tempfile.mkdtemp(
-            suffix="charm-upgrade", prefix="tmp")
-
-    def retrieve_charm(self, charm_id):
-        return download_charm(
-            self._agent.client, charm_id, self._charm_directory)
-
-    def _remove_tree(self, result):
-        if os.path.exists(self._charm_directory):
-            shutil.rmtree(self._charm_directory)
-        return result
-
-    def run(self):
-        d = self._run()
-        d.addBoth(self._remove_tree)
-        return d
-
-    @inlineCallbacks
-    def _run(self):
-        self._log.info("Starting charm upgrade...")
-
-        # Verify the workflow state
-        workflow_state = yield self._agent.workflow.get_state()
-        if workflow_state != "started":
-            self._log.warning(
-                "Unit not in an upgradeable state: %s", workflow_state)
-            # Upgrades can only be supported while the unit is
-            # running, we clear the flag because we don't support
-            # persistent upgrade requests across unit starts. The
-            # upgrade request will need to be reissued, after
-            # resolving or restarting the unit.
-            yield self._agent.unit_state.clear_upgrade_flag()
-            returnValue(False)
-
-        # Get, check, and clear the flag. Do it first so a second upgrade
-        # will restablish the upgrade request.
-        upgrade_flag = yield self._agent.unit_state.get_upgrade_flag()
-        if not upgrade_flag:
-            self._log.warning("No upgrade flag set.")
-            returnValue(False)
-
-        self._log.debug("Clearing upgrade flag.")
-        yield self._agent.unit_state.clear_upgrade_flag()
-
-        # Retrieve the service state
-        service_state_manager = ServiceStateManager(self._agent.client)
-        service_state = yield service_state_manager.get_service_state(
-            self._agent.unit_name.split("/")[0])
-
-        # Verify unit state, upgrade flag, and newer version requested.
-        service_charm_id = yield service_state.get_charm_id()
-        unit_charm_id = yield self._agent.unit_state.get_charm_id()
-
-        if service_charm_id == unit_charm_id:
-            self._log.debug("Unit already running latest charm")
-            yield self._agent.unit_state.clear_upgrade_flag()
-            returnValue(True)
-
-        # Retrieve charm
-        self._log.debug("Retrieving charm %s", service_charm_id)
-        charm = yield self.retrieve_charm(service_charm_id)
-
-        # Stop hook executions
-        self._log.debug("Stopping hook execution.")
-        yield self._agent.executor.stop()
-
-        # Note the current charm version
-        self._log.debug("Setting unit charm id to %s", service_charm_id)
-        yield self._agent.unit_state.set_charm_id(service_charm_id)
-
-        # Extract charm
-        self._log.debug("Extracting new charm.")
-        charm.extract_to(
-            os.path.join(self._agent.unit_directory, "charm"))
-
-        # Upgrade
-        self._log.debug("Invoking upgrade transition.")
-
-        success = yield self._agent.workflow.fire_transition(
-            "upgrade_charm")
-
-        if success:
-            self._log.debug("Unit upgraded.")
-        else:
-            self._log.warning("Upgrade failed.")
-
-        returnValue(success)
-
 if __name__ == '__main__':
     UnitAgent.run()

=== modified file 'juju/control/options.py'
--- juju/control/options.py	2011-01-20 18:00:23 +0000
+++ juju/control/options.py	2012-02-02 16:42:42 +0000
@@ -53,9 +53,8 @@
         )
 
     unix_group.add_argument(
-        "--pidfile", default="%s.pid" % agent.name,
+        "--pidfile", default="",
         help="Path to the pid file",
-        type=ensure_abs_path,
         )
 
     unix_group.add_argument(
@@ -91,9 +90,9 @@
         )
 
     unix_group.add_argument(
-        "--daemon", "-n", default=True,
-        dest="nodaemon", action="store_false",
-        help="Daemonize the process",
+        "--nodaemon", "-n", default=False,
+        dest="nodaemon", action="store_true",
+        help="Don't daemonize (stay in foreground)",
         )
 
     unix_group.add_argument(

=== modified file 'juju/control/status.py'
--- juju/control/status.py	2011-12-07 05:02:08 +0000
+++ juju/control/status.py	2012-02-02 16:42:42 +0000
@@ -216,8 +216,10 @@
             relation_status = {}
             for relation in relations:
                 try:
+                    print unit.unit_name
                     relation_unit = yield relation.get_unit_state(unit)
                 except UnitRelationStateNotFound:
+                    print "POW SPLAT"
                     # This exception will occur when relations are
                     # established between services without service
                     # units, and therefore never have any

=== modified file 'juju/control/tests/test_resolved.py'
--- juju/control/tests/test_resolved.py	2012-01-12 10:18:07 +0000
+++ juju/control/tests/test_resolved.py	2012-02-02 16:42:42 +0000
@@ -88,10 +88,10 @@
         """
         for unit, state in units:
             unit_relation = yield service_relation.add_unit_state(unit)
-            lifecycle = UnitRelationLifecycle(self.client,
-                                              unit.unit_name, unit_relation,
-                                              service_relation.relation_name,
-                                              self.makeDir(), self.executor)
+            lifecycle = UnitRelationLifecycle(
+                self.client, unit.unit_name, unit_relation,
+                service_relation.relation_name, self.makeDir(), self.makeDir(),
+                self.executor)
             workflow_state = RelationWorkflowState(
                 self.client, unit_relation, service_relation.relation_name,
                 lifecycle, self.makeDir())

=== modified file 'juju/control/tests/test_status.py'
--- juju/control/tests/test_status.py	2011-12-07 18:29:12 +0000
+++ juju/control/tests/test_status.py	2012-02-02 16:42:42 +0000
@@ -39,7 +39,7 @@
     # Status tests setup a large tree every time, make allowances for it.
     # TODO: create minimal trees needed per test.
     timeout = 10
-    
+
     @inlineCallbacks
     def setUp(self):
         yield super(StatusTestBase, self).setUp()
@@ -107,6 +107,7 @@
         options = TwistedOptionNamespace()
         options["juju_directory"] = path
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         for k, v in extra_options.items():
             options[k] = v
         agent.configure(options)
@@ -302,6 +303,7 @@
         options = TwistedOptionNamespace()
         options["juju_directory"] = self.makeDir()
         options["zookeeper_servers"] = get_test_zookeeper_address()
+        options["session_file"] = self.makeFile()
         options["machine_id"] = "0"
         agent.configure(options)
         agent.set_watch_enabled(False)

=== modified file 'juju/errors.py'
--- juju/errors.py	2011-09-24 22:21:23 +0000
+++ juju/errors.py	2012-02-02 16:42:42 +0000
@@ -62,7 +62,7 @@
         return "Error processing %r: %s" % (self.path, self.message)
 
 
-class CharmInvocationError(JujuError):
+class CharmInvocationError(CharmError):
     """A charm's hook invocation exited with an error"""
 
     def __init__(self, path, exit_code):
@@ -74,6 +74,16 @@
             self.path, self.exit_code)
 
 
+class CharmUpgradeError(CharmError):
+    """Something went wrong trying to upgrade a charm"""
+
+    def __init__(self, message):
+        self.message = message
+
+    def __str__(self):
+        return "Cannot upgrade charm: %s" % self.message
+
+
 class FileAlreadyExists(JujuError):
     """Raised when something refuses to overwrite an existing file.
 
@@ -164,3 +174,6 @@
                 self.user_policy,
                 self.provider_type,
                 ", ".join(self.provider_policies)))
+
+class ServiceError(JujuError):
+    """Some problem with an upstart service"""

=== modified file 'juju/hooks/scheduler.py'
--- juju/hooks/scheduler.py	2011-12-12 01:56:05 +0000
+++ juju/hooks/scheduler.py	2012-02-02 16:42:42 +0000
@@ -1,4 +1,6 @@
 import logging
+import os
+import yaml
 
 from twisted.internet.defer import DeferredQueue, inlineCallbacks
 from juju.state.hook import RelationHookContext, RelationChange
@@ -28,31 +30,69 @@
     the run queue.
     """
 
-    def __init__(self, client, executor, unit_relation, relation_name, unit_name):
-        self._running = None
+    def __init__(self, client, executor, unit_relation, relation_name,
+                 unit_name, state_path):
+        self._running = False
+        self._state_path = state_path
 
         # The thing that will actually run the hook for us
         self._executor = executor
-
         # For hook context construction.
         self._client = client
         self._unit_relation = unit_relation
         self._relation_name = relation_name
-        self._members = None
         self._unit_name = unit_name
 
-        # Track next operation by node
-        self._node_queue = {}
-
-        # Track node operations by clock tick
-        self._clock_queue = {}
-
+        if os.path.exists(self._state_path):
+            self._load_state()
+        else:
+            self._create_state()
+
+    def _create_state(self):
+        # Current units (as far as the next hook should know)
+        self._context_members = None
+        # Current units and settings versions (as far as the queue knows)
+        self._member_versions = {}
+        # Tracks next operation by unit
+        self._unit_ops = {}
+        # Tracks unit operations by clock tick
+        self._clock_units = {}
         # Run queue (clock)
         self._run_queue = DeferredQueue()
-
         # Artifical clock sequence
         self._clock_sequence = 0
 
+    def _load_state(self):
+        with open(self._state_path) as f:
+            state = yaml.load(f.read())
+            if not state:
+                return self._create_state()
+        self._context_members = state["context_members"]
+        self._member_versions = state["member_versions"]
+        self._unit_ops = state["unit_ops"]
+        self._clock_units = state["clock_units"]
+        self._run_queue = DeferredQueue()
+        self._run_queue.pending = state["clock_queue"]
+        self._clock_sequence = state["clock_sequence"]
+
+    def _save_state(self):
+        state = yaml.dump({
+            "context_members": self._context_members,
+            "member_versions": self._member_versions,
+            "unit_ops": self._unit_ops,
+            "clock_units": self._clock_units,
+            "clock_queue": [
+                # Strip "stop" instructions: if the lifecycle stopped us,
+                # then if/when the lifecycle comes up again in a stopped
+                # state, it won't start us in the first place.
+                c for c in self._run_queue.pending if c is not None],
+            "clock_sequence": self._clock_sequence})
+
+        temp_path = self._state_path + "~"
+        with open(temp_path, "w") as f:
+            f.write(state)
+        os.rename(temp_path, self._state_path)
+
     @property
     def running(self):
         return self._running is True
@@ -61,6 +101,12 @@
     def run(self):
         """Run the hook scheduler and execution."""
         assert not self._running, "Scheduler is already running"
+        try:
+            with open(self._state_path, "a"):
+                pass
+        except IOError:
+            raise AssertionError("%s is not writable!" % self._state_path)
+
         self._running = True
         log.debug("start")
 
@@ -72,16 +118,17 @@
                 break
 
             # Get all the units with changes in this clock tick.
-            for unit_name in self._clock_queue.pop(clock):
+            for unit_name in self._clock_units.pop(clock):
 
                 # Get the change for the unit.
-                change_clock, change_type = self._node_queue.pop(unit_name)
+                change_clock, change_type = self._unit_ops.pop(unit_name)
 
                 log.debug("executing hook for %s:%s",
                           unit_name, CHANGE_LABELS[change_type])
 
                 # Execute the hook
                 yield self._execute(unit_name, change_type)
+            self._save_state()
 
     def stop(self):
         """Stop the hook execution.
@@ -99,26 +146,24 @@
         # occurs.
         self._run_queue.put(None)
 
-    def notify_change(self, old_units=(), new_units=(), modified=()):
-        """Receive changes regarding related units and schedule hook execution.
-        """
-        log.debug("relation change old:%s, new:%s, modified:%s",
-                  old_units, new_units, modified)
-
+    def cb_change_members(self, old_units, new_units):
+        log.debug("members changed: old=%s, new=%s", old_units, new_units)
+        scheduled = 0
         self._clock_sequence += 1
 
-        # keep track if we've scheduled changes during this clock
-        scheduled = 0
-
-        # Handle membership changes
-
-        # If we don't have a cached membership yet, use the old units
-        # as a baseline.
-        if self._members is None:
-            self._members = list(old_units)
-
-        added = set(new_units) - set(old_units)
-        removed = set(old_units) - set(new_units)
+        if self._context_members is None:
+            self._context_members = list(old_units)
+
+        if set(self._member_versions) != set(old_units):
+            log.debug(
+                "old does not match last recorded units: %s",
+                sorted(self._member_versions))
+
+        added = set(new_units) - set(self._member_versions)
+        removed = set(self._member_versions) - set(new_units)
+        self._member_versions.update(dict((unit, 0) for unit in added))
+        for unit in removed:
+            del self._member_versions[unit]
 
         for unit_name in sorted(added):
             scheduled += self._queue_change(
@@ -128,59 +173,68 @@
             scheduled += self._queue_change(
                 unit_name, REMOVED, self._clock_sequence)
 
-        # Handle modified change
-        for unit_name in modified:
-            scheduled += self._queue_change(
-                unit_name, MODIFIED, self._clock_sequence)
+        if scheduled:
+            self._run_queue.put(self._clock_sequence)
+        self._save_state()
 
+    def cb_change_settings(self, unit_versions):
+        log.debug("settings changed: %s", unit_versions)
+        scheduled = 0
+        self._clock_sequence += 1
+        for (unit_name, version) in unit_versions:
+            if version > self._member_versions.get(unit_name, 0):
+                self._member_versions[unit_name] = version
+                scheduled += self._queue_change(
+                    unit_name, MODIFIED, self._clock_sequence)
         if scheduled:
             self._run_queue.put(self._clock_sequence)
+        self._save_state()
 
     def get_hook_context(self, change):
         """
         Return a hook context, corresponding to the current state of the
         system.
         """
-        members = self._members or ()
+        context_members = self._context_members or ()
         context = RelationHookContext(
             self._client, self._unit_relation, change,
-            sorted(members), unit_name=self._unit_name)
+            sorted(context_members), unit_name=self._unit_name)
         return context
 
     def _queue_change(self, unit_name, change_type, clock):
         """Queue up the node change for execution.
         """
         # If its a new change for the unit store it, and return.
-        if not unit_name in self._node_queue:
-            self._node_queue[unit_name] = (clock, change_type)
-            self._clock_queue.setdefault(clock, []).append(unit_name)
+        if not unit_name in self._unit_ops:
+            self._unit_ops[unit_name] = (clock, change_type)
+            self._clock_units.setdefault(clock, []).append(unit_name)
             return True
 
         # Else merge/reduce with the previous operation.
-        previous_clock, previous_change = self._node_queue[unit_name]
+        previous_clock, previous_change = self._unit_ops[unit_name]
         change_clock, change_type = self._reduce(
             (previous_clock, previous_change),
             (self._clock_sequence, change_type))
 
         # If they've cancelled, remove from node and clock queues
         if change_type is None:
-            del self._node_queue[unit_name]
-            self._clock_queue[previous_clock].remove(unit_name)
+            del self._unit_ops[unit_name]
+            self._clock_units[previous_clock].remove(unit_name)
             return False
 
         # Update the node queue with the merged change.
-        self._node_queue[unit_name] = (change_clock, change_type)
+        self._unit_ops[unit_name] = (change_clock, change_type)
 
         # If the clock has changed, remove the old entry.
         if change_clock != previous_clock:
-            self._clock_queue[previous_clock].remove(unit_name)
+            self._clock_units[previous_clock].remove(unit_name)
 
         # If the old entry has precedence, we didn't schedule anything for
         # this clock cycle.
         if change_clock != clock:
             return False
 
-        self._clock_queue.setdefault(clock, []).append(unit_name)
+        self._clock_units.setdefault(clock, []).append(unit_name)
         return True
 
     def _reduce(self, previous, new):
@@ -214,9 +268,9 @@
         """
         # Determine the current members as of the change.
         if change_type == ADDED:
-            self._members.append(unit_name)
+            self._context_members.append(unit_name)
         elif change_type == REMOVED:
-            self._members.remove(unit_name)
+            self._context_members.remove(unit_name)
 
         # Assemble the change and hook execution context
         change = RelationChange(

=== modified file 'juju/hooks/tests/test_scheduler.py'
--- juju/hooks/tests/test_scheduler.py	2011-12-12 01:56:05 +0000
+++ juju/hooks/tests/test_scheduler.py	2012-02-02 16:42:42 +0000
@@ -1,11 +1,18 @@
 import logging
-
-from twisted.internet.defer import inlineCallbacks
-
-from juju.hooks.scheduler import HookScheduler
+import os
+import yaml
+
+from twisted.internet.defer import inlineCallbacks, fail, succeed
+
+from juju.hooks.scheduler import HookScheduler, ADDED, MODIFIED, REMOVED
 from juju.state.hook import RelationChange
 from juju.state.tests.test_service import ServiceStateManagerTestBase
 
+
+class SomeError(Exception):
+    pass
+
+
 class HookSchedulerTest(ServiceStateManagerTestBase):
 
     @inlineCallbacks
@@ -15,29 +22,49 @@
         self.unit_relation = self.mocker.mock()
         self.executions = []
         self.service = yield self.add_service_from_charm("wordpress")
-        self.scheduler = HookScheduler(self.client,
-                                       self.collect_executor,
-                                       self.unit_relation, "",
-                                       unit_name="wordpress/0")
+        self.state_file = self.makeFile()
+        self.executor = self.collect_executor
+        self._scheduler = None
         self.log_stream = self.capture_logging(
             "hook.scheduler", level=logging.DEBUG)
 
+    @property
+    def scheduler(self):
+        # Create lazily, so we can create with a state file if we want to,
+        # and swap out collect_executor when helpful to do so.
+        if self._scheduler is None:
+            self._scheduler = HookScheduler(
+                self.client, self.executor, self.unit_relation, "",
+                "wordpress/0", self.state_file)
+        return self._scheduler
+
     def collect_executor(self, context, change):
         self.executions.append((context, change))
 
+    def write_single_unit_state(self):
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({
+                "context_members": ["u-1"],
+                "member_versions": {"u-1": 0},
+                "unit_ops": {},
+                "clock_units": {},
+                "clock_queue": [],
+                "clock_sequence": 1}))
+
     # Event reduction/coalescing cases
     def test_reduce_removed_added(self):
         """ A remove event for a node followed by an add event,
         results in a modify event.
         """
-        self.scheduler.notify_change(old_units=["u-1"], new_units=[])
-        self.scheduler.notify_change(old_units=[], new_units=["u-1"])
+        self.write_single_unit_state()
+        self.scheduler.cb_change_members(["u-1"], [])
+        self.scheduler.cb_change_members([], ["u-1"])
         self.scheduler.run()
         self.assertEqual(len(self.executions), 1)
         self.assertEqual(self.executions[0][1].change_type, "modified")
 
-        output = ("relation change old:['u-1'], new:[], modified:()",
-                  "relation change old:[], new:['u-1'], modified:()",
+        output = ("members changed: old=['u-1'], new=[]",
+                  "members changed: old=[], new=['u-1']",
                   "start",
                   "executing hook for u-1:modified\n")
         self.assertEqual(self.log_stream.getvalue(), "\n".join(output))
@@ -46,34 +73,34 @@
         """A modify, remove, add event for a node results in a modify.
         An extra validation of the previous test.
         """
-        self.scheduler.notify_change(modified=["u-1"])
-        self.scheduler.notify_change(old_units=["u-1"], new_units=[])
-        self.scheduler.notify_change(old_units=[], new_units=["u-1"])
+        self.write_single_unit_state()
+        self.scheduler.cb_change_settings([("u-1", 1)])
+        self.scheduler.cb_change_members(["u-1"], [])
+        self.scheduler.cb_change_members([], ["u-1"])
         self.scheduler.run()
         self.assertEqual(len(self.executions), 1)
         self.assertEqual(self.executions[0][1].change_type, "modified")
 
     def test_reduce_add_modify(self):
         """An add and modify event for a node are coalesced to an add."""
-        self.scheduler.notify_change(old_units=[], new_units=["u-1"])
-        self.scheduler.notify_change(modified=["u-1"])
+        self.scheduler.cb_change_members([], ["u-1"])
+        self.scheduler.cb_change_settings([("u-1", 1)])
         self.scheduler.run()
         self.assertEqual(len(self.executions), 1)
         self.assertEqual(self.executions[0][1].change_type, "joined")
 
     def test_reduce_add_remove(self):
         """an add followed by a removal results in a noop."""
-        self.scheduler.notify_change(old_units=[], new_units=["u-1"])
-        self.scheduler.notify_change(old_units=["u-1"], new_units=[])
+        self.scheduler.cb_change_members([], ["u-1"])
+        self.scheduler.cb_change_members(["u-1"], [])
         self.scheduler.run()
         self.assertEqual(len(self.executions), 0)
 
     def test_reduce_modify_remove(self):
         """Modifying and then removing a node, results in just the removal."""
-        self.scheduler.notify_change(old_units=["u-1"],
-                                     new_units=["u-1"],
-                                     modified=["u-1"])
-        self.scheduler.notify_change(old_units=["u-1"], new_units=[])
+        self.write_single_unit_state()
+        self.scheduler.cb_change_settings([("u-1", 1)])
+        self.scheduler.cb_change_members(["u-1"], [])
         self.scheduler.run()
         self.assertEqual(len(self.executions), 1)
         self.assertEqual(self.executions[0][1].change_type, "departed")
@@ -82,15 +109,15 @@
         """Multiple modifies get coalesced to a single modify."""
         # simulate normal startup, the first notify will always be the existing
         # membership set.
-        self.scheduler.notify_change(old_units=[], new_units=["u-1"])
+        self.scheduler.cb_change_members([], ["u-1"])
         self.scheduler.run()
         self.scheduler.stop()
         self.assertEqual(len(self.executions), 1)
 
         # Now continue the modify/modify reduction.
-        self.scheduler.notify_change(modified=["u-1"])
-        self.scheduler.notify_change(modified=["u-1"])
-        self.scheduler.notify_change(modified=["u-1"])
+        self.scheduler.cb_change_settings([("u-1", 1)])
+        self.scheduler.cb_change_settings([("u-1", 2)])
+        self.scheduler.cb_change_settings([("u-1", 3)])
         self.scheduler.run()
 
         self.assertEqual(len(self.executions), 2)
@@ -112,20 +139,35 @@
         self.assertFalse(self.scheduler.running)
 
     @inlineCallbacks
+    def test_run_requires_writable_state(self):
+        # Induce lazy creation of scheduler, then break state file
+        self.scheduler
+        with open(self.state_file, "w"):
+            pass
+        os.chmod(self.state_file, 0)
+        e = yield self.assertFailure(self.scheduler.run(), AssertionError)
+        self.assertEquals(str(e), "%s is not writable!" % self.state_file)
+
+    def test_empty_state(self):
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({}))
+
+        # Induce lazy creation to verify it can still survive
+        self.scheduler
+
+    @inlineCallbacks
     def test_membership_visibility_per_change(self):
         """Hooks are executed against changes, those changes are
         associated to a temporal timestamp, however the changes
         are scheduled for execution, and the state/time of the
         world may have advanced, to present a logically consistent
-        view, we try to gaurantee at a minimum, that hooks will
-        always see the membership of a relations it was at the
+        view, we try to guarantee at a minimum, that hooks will
+        always see the membership of a relation as it was at the
         time of their associated change.
         """
-        self.scheduler.notify_change(
-            old_units=[], new_units=["u-1", "u-2"])
-        self.scheduler.notify_change(
-            old_units=["u-1", "u-2"], new_units=["u-2", "u-3"])
-        self.scheduler.notify_change(modified=["u-2"])
+        self.scheduler.cb_change_members([], ["u-1", "u-2"])
+        self.scheduler.cb_change_members(["u-1", "u-2"], ["u-2", "u-3"])
+        self.scheduler.cb_change_settings([("u-2", 1)])
 
         self.scheduler.run()
         self.scheduler.stop()
@@ -139,9 +181,8 @@
         change_members = yield self.executions[0][0].get_members()
         self.assertEqual(change_members, ["u-2"])
 
-        self.scheduler.notify_change(modified=["u-2"])
-        self.scheduler.notify_change(
-            old_units=["u-2", "u-3"], new_units=["u-2"])
+        self.scheduler.cb_change_settings([("u-2", 2)])
+        self.scheduler.cb_change_members(["u-2", "u-3"], ["u-2"])
         self.scheduler.run()
 
         self.assertEqual(len(self.executions), 4)
@@ -156,10 +197,17 @@
         a hook wont see any 'active' members in a membership list, that
         it hasn't previously been given a notify of before.
         """
-        self.scheduler.notify_change(
-            old_units=["u-1", "u-2"],
-            new_units=["u-2", "u-3", "u-4"],
-            modified=["u-2"])
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({
+                "context_members": ["u-1", "u-2"],
+                "member_versions": {"u-1": 0, "u-2": 0},
+                "unit_ops": {},
+                "clock_units": {},
+                "clock_queue": [],
+                "clock_sequence": 1}))
+
+        self.scheduler.cb_change_members(["u-1", "u-2"], ["u-2", "u-3", "u-4"])
+        self.scheduler.cb_change_settings([("u-2", 1)])
 
         self.scheduler.run()
         self.scheduler.stop()
@@ -197,3 +245,172 @@
             RelationChange("", "", ""))
         members = yield context.get_members()
         self.assertEqual(members, [])
+
+    @inlineCallbacks
+    def test_state_is_loaded(self):
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({
+                "context_members": ["u-1", "u-2"],
+                "member_versions": {"u-1": 5, "u-2": 2, "u-3": 0},
+                "unit_ops": {"u-1": (3, MODIFIED), "u-3": (4, ADDED)},
+                "clock_units": {3: ["u-1"], 4: ["u-3"]},
+                "clock_queue": [3, 4],
+                "clock_sequence": 4}))
+
+        self.scheduler.run()
+        while len(self.executions) < 2:
+            yield self.poke_zk()
+        self.scheduler.stop()
+
+        self.assertEqual(self.executions[0][1].change_type, "modified")
+        members = yield self.executions[0][0].get_members()
+        self.assertEqual(members, ["u-1", "u-2"])
+
+        self.assertEqual(self.executions[1][1].change_type, "joined")
+        members = yield self.executions[1][0].get_members()
+        self.assertEqual(members, ["u-1", "u-2", "u-3"])
+
+        with open(self.state_file) as f:
+            state = yaml.load(f.read())
+        self.assertEquals(state, {
+            "context_members": ["u-1", "u-2", "u-3"],
+            "member_versions": {"u-1": 5, "u-2": 2, "u-3": 0},
+            "unit_ops": {},
+            "clock_units": {},
+            "clock_queue": [],
+            "clock_sequence": 4})
+
+    def test_state_is_stored(self):
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({
+                "context_members": ["u-1", "u-2"],
+                "member_versions": {"u-1": 0, "u-2": 2},
+                "unit_ops": {},
+                "clock_units": {},
+                "clock_queue": [],
+                "clock_sequence": 7}))
+
+        self.scheduler.cb_change_members(["u-1", "u-2"], ["u-2", "u-3"])
+        self.scheduler.cb_change_settings([("u-2", 3)])
+
+        # Add a stop instruction to the queue, which should *not* be saved.
+        self.scheduler.stop()
+
+        with open(self.state_file) as f:
+            state = yaml.load(f.read())
+        self.assertEquals(state, {
+            "context_members": ["u-1", "u-2"],
+            "member_versions": {"u-2": 3, "u-3": 0},
+            "unit_ops": {"u-1": (8, REMOVED),
+                         "u-2": (9, MODIFIED),
+                         "u-3": (8, ADDED)},
+            "clock_units": {8: ["u-3", "u-1"], 9: ["u-2"]},
+            "clock_queue": [8, 9],
+            "clock_sequence": 9})
+
+    @inlineCallbacks
+    def test_state_stored_after_tick(self):
+
+        def execute(context, change):
+            self.execute_calls += 1
+            if self.execute_calls > 1:
+                return fail(SomeError())
+            return succeed(None)
+        self.execute_calls = 0
+        self.executor = execute
+
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({
+            "context_members": ["u-1", "u-2"],
+            "member_versions": {"u-1": 1, "u-2": 0, "u-3": 0},
+            "unit_ops": {"u-1": (3, MODIFIED), "u-3": (4, ADDED)},
+            "clock_units": {3: ["u-1"], 4: ["u-3"]},
+            "clock_queue": [3, 4],
+            "clock_sequence": 4}))
+
+        d = self.scheduler.run()
+        while self.execute_calls < 2:
+            yield self.poke_zk()
+        yield self.assertFailure(d, SomeError)
+        with open(self.state_file) as f:
+            self.assertEquals(yaml.load(f.read()), {
+                "context_members": ["u-1", "u-2"],
+                "member_versions": {"u-1": 1, "u-2": 0, "u-3": 0},
+                "unit_ops": {"u-3": (4, ADDED)},
+                "clock_units": {4: ["u-3"]},
+                "clock_queue": [4],
+                "clock_sequence": 4})
+
+    @inlineCallbacks
+    def test_state_not_stored_mid_tick(self):
+
+        def execute(context, change):
+            self.execute_called = True
+            return fail(SomeError())
+        self.execute_called = False
+        self.executor = execute
+
+        initial_state = {
+            "context_members": ["u-1", "u-2"],
+            "member_versions": {"u-1": 1, "u-2": 0, "u-3": 0},
+            "unit_ops": {"u-1": (3, MODIFIED), "u-3": (4, ADDED)},
+            "clock_units": {3: ["u-1"], 4: ["u-3"]},
+            "clock_queue": [3, 4],
+            "clock_sequence": 4}
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump(initial_state))
+
+        d = self.scheduler.run()
+        while not self.execute_called:
+            yield self.poke_zk()
+        yield self.assertFailure(d, SomeError)
+        with open(self.state_file) as f:
+            self.assertEquals(yaml.load(f.read()), initial_state)
+
+    def test_ignore_equal_settings_version(self):
+        """
+        A modified event whose version is not greater than the latest known
+        version for that unit will be ignored.
+        """
+        self.write_single_unit_state()
+        self.scheduler.cb_change_settings([("u-1", 0),])
+        self.scheduler.run()
+        self.assertEquals(len(self.executions), 0)
+
+    def test_settings_version_0_on_add(self):
+        """
+        When a unit is added, we assume its settings version to be 0, and
+        therefore modified events with version 0 will be ignored.
+        """
+        self.scheduler.cb_change_members([], ["u-1"])
+        self.scheduler.cb_change_settings([("u-1", 0),])
+        self.scheduler.run()
+        self.assertEquals(len(self.executions), 1)
+        self.assertEqual(self.executions[0][1].change_type, "joined")
+
+    def test_membership_timeslip(self):
+        """
+        Adds and removes are calculated based on known membership state, NOT
+        on old_units.
+        """
+        with open(self.state_file, "w") as f:
+            f.write(yaml.dump({
+            "context_members": ["u-1", "u-2"],
+            "member_versions": {"u-1": 0, "u-2": 0},
+            "unit_ops": {},
+            "clock_units": {},
+            "clock_queue": [],
+            "clock_sequence": 4}))
+
+        self.scheduler.cb_change_members(["u-2"], ["u-3", "u-4"])
+        self.scheduler.run()
+
+        output = (
+            "members changed: old=['u-2'], new=['u-3', 'u-4']",
+            "old does not match last recorded units: ['u-1', 'u-2']",
+            "start",
+            "executing hook for u-3:joined",
+            "executing hook for u-4:joined",
+            "executing hook for u-1:departed",
+            "executing hook for u-2:departed\n")
+        self.assertEqual(self.log_stream.getvalue(), "\n".join(output))

=== modified file 'juju/lib/lxc/tests/test_lxc.py'
--- juju/lib/lxc/tests/test_lxc.py	2011-10-01 00:04:14 +0000
+++ juju/lib/lxc/tests/test_lxc.py	2012-02-02 16:42:42 +0000
@@ -10,10 +10,17 @@
 from juju.lib.testing import TestCase
 
 
-def run_lxc_tests():
-    if os.environ.get("TEST_LXC"):
-        return None
-    return "TEST_LXC=1 to include lxc tests"
+def skip_sudo_tests():
+    if os.environ.get("TEST_SUDO"):
+        # Get user's password *now*, if needed, not mid-run
+        os.system("sudo false")
+        return False
+    return "TEST_SUDO=1 to include tests which use sudo (including lxc tests)"
+
+
+def uses_sudo(f):
+    f.skip = skip_sudo_tests()
+    return f
 
 
 DATA_PATH = os.path.abspath(
@@ -23,9 +30,9 @@
 DEFAULT_CONTAINER = "lxc_test"
 
 
+@uses_sudo
 class LXCTest(TestCase):
     timeout = 240
-    skip = run_lxc_tests()
 
     def setUp(self):
         self.config = self.make_config()

=== added directory 'juju/lib/tests/data'
=== added file 'juju/lib/tests/data/test_basic_install'
--- juju/lib/tests/data/test_basic_install	1970-01-01 00:00:00 +0000
+++ juju/lib/tests/data/test_basic_install	2012-02-02 16:42:42 +0000
@@ -0,0 +1,10 @@
+description "uninteresting service"
+author "Juju Team <[e-mail]>"
+
+start on runlevel [2345]
+stop on runlevel [!2345]
+respawn
+
+
+
+exec /bin/false >> /tmp/some-name.output 2>&1

=== added file 'juju/lib/tests/data/test_less_basic_install'
--- juju/lib/tests/data/test_less_basic_install	1970-01-01 00:00:00 +0000
+++ juju/lib/tests/data/test_less_basic_install	2012-02-02 16:42:42 +0000
@@ -0,0 +1,11 @@
+description "pew pew pew blam"
+author "Juju Team <[e-mail]>"
+
+start on runlevel [2345]
+stop on runlevel [!2345]
+respawn
+
+env FOO="bar baz qux"
+env PEW="pew"
+
+exec /bin/deathstar --ignore-ewoks endor >> /somewhere/else 2>&1

=== added file 'juju/lib/tests/data/test_standard_install'
--- juju/lib/tests/data/test_standard_install	1970-01-01 00:00:00 +0000
+++ juju/lib/tests/data/test_standard_install	2012-02-02 16:42:42 +0000
@@ -0,0 +1,10 @@
+description "a wretched hive of scum and villainy"
+author "Juju Team <[e-mail]>"
+
+start on runlevel [2345]
+stop on runlevel [!2345]
+respawn
+
+env LIGHTSABER="civilised weapon"
+
+exec /bin/imagination-failure --no-ideas >> /tmp/some-name.output 2>&1

=== modified file 'juju/lib/tests/test_statemachine.py'
--- juju/lib/tests/test_statemachine.py	2011-09-15 18:50:23 +0000
+++ juju/lib/tests/test_statemachine.py	2012-02-02 16:42:42 +0000
@@ -121,10 +121,14 @@
         workflow_state = AttributeWorkflowState(workflow)
         current_state = yield workflow_state.get_state()
         self.assertEqual(current_state, None)
+        current_vars = yield workflow_state.get_state_variables()
+        self.assertEqual(current_vars, {})
 
         yield workflow_state.set_state("started")
         current_state = yield workflow_state.get_state()
         self.assertEqual(current_state, "started")
+        current_vars = yield workflow_state.get_state_variables()
+        self.assertEqual(current_vars, {})
 
     @inlineCallbacks
     def test_state_fire_transition(self):
@@ -333,3 +337,13 @@
 
         self.assertFailure(workflow_state.transition_state("unknown"),
                            InvalidStateError)
+
+    @inlineCallbacks
+    def test_load_bad_state(self):
+        class BadLoadWorkflowState(WorkflowState):
+            def _load(self):
+                return succeed({"some": "other-data"})
+
+        workflow = BadLoadWorkflowState(Workflow())
+        yield self.assertFailure(workflow.get_state(), KeyError)
+        yield self.assertFailure(workflow.get_state_variables(), KeyError)

=== added file 'juju/lib/tests/test_upstart.py'
--- juju/lib/tests/test_upstart.py	1970-01-01 00:00:00 +0000
+++ juju/lib/tests/test_upstart.py	2012-02-02 16:42:42 +0000
@@ -0,0 +1,339 @@
+import os
+
+from twisted.internet.defer import inlineCallbacks, succeed
+
+from juju.errors import ServiceError
+from juju.lib.mocker import ANY, KWARGS
+from juju.lib.testing import TestCase
+from juju.lib.upstart import UpstartService
+
+
+DATA_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "data")
+
+
+class UpstartServiceTest(TestCase):
+
+    @inlineCallbacks
+    def setUp(self):
+        yield super(UpstartServiceTest, self).setUp()
+        self.init_dir = self.makeDir()
+        self.conf = os.path.join(self.init_dir, "some-name.conf")
+        self.output = "/tmp/some-name.output"
+        self.patch(UpstartService, "init_dir", self.init_dir)
+        self.service = UpstartService("some-name")
+
+    def setup_service(self):
+        self.service.set_description("a wretched hive of scum and villainy")
+        self.service.set_command("/bin/imagination-failure --no-ideas")
+        self.service.set_environ({"LIGHTSABER": "civilised weapon"})
+
+    def setup_mock(self):
+        self.check_call = self.mocker.replace("subprocess.check_call")
+        self.getProcessOutput = self.mocker.replace(
+            "twisted.internet.utils.getProcessOutput")
+
+    def mock_status(self, result):
+        self.getProcessOutput("/sbin/status", ["some-name"])
+        self.mocker.result(result)
+
+    def mock_call(self, args, output=None):
+        self.check_call(args, KWARGS)
+        if output is None:
+            self.mocker.result(0)
+        else:
+            def write(ANY, **_):
+                with open(self.output, "w") as f:
+                    f.write(output)
+            self.mocker.call(write)
+
+    def mock_start(self, output=None):
+        self.mock_call(("/sbin/start", "some-name"), output)
+
+    def mock_stop(self):
+        self.mock_call(("/sbin/stop", "some-name"))
+
+    def mock_check_success(self):
+        for _ in range(5):
+            self.mock_status(succeed("blah start/running blah 12345"))
+
+    def mock_check_unstable(self):
+        for _ in range(4):
+            self.mock_status(succeed("blah start/running blah 12345"))
+        self.mock_status(succeed("blah start/running blah 12346"))
+
+    def mock_check_not_running(self):
+        self.mock_status(succeed("blah"))
+
+    def write_dummy_conf(self):
+        with open(self.conf, "w") as f:
+            f.write("dummy")
+
+    def assert_dummy_conf(self):
+        with open(self.conf) as f:
+            self.assertEquals(f.read(), "dummy")
+
+    def assert_no_conf(self):
+        self.assertFalse(os.path.exists(self.conf))
+
+    def assert_conf(self, name="test_standard_install"):
+        with open(os.path.join(DATA_DIR, name)) as expected:
+            with open(self.conf) as actual:
+                self.assertEquals(actual.read(), expected.read())
+
+    def test_is_installed(self):
+        """Check is_installed depends on conf file existence"""
+        self.assertFalse(self.service.is_installed())
+        self.write_dummy_conf()
+        self.assertTrue(self.service.is_installed())
+
+    def test_init_dir(self):
+        """
+        Check is_installed still works when init_dir specified explicitly
+        """
+        self.patch(UpstartService, "init_dir", "/BAD/PATH")
+        self.service = UpstartService("some-name", init_dir=self.init_dir)
+        self.setup_service()
+
+        self.assertFalse(self.service.is_installed())
+        self.write_dummy_conf()
+        self.assertTrue(self.service.is_installed())
+
+    @inlineCallbacks
+    def test_is_running(self):
+        """
+        Check is_running interprets status output (when service is installed)
+        """
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_status(succeed("blah blob/gibbering blah"))
+        self.mock_status(succeed("blah start/running blah 12345"))
+        self.mocker.replay()
+
+        # Won't hit status; conf is not installed
+        self.assertFalse((yield self.service.is_running()))
+        self.write_dummy_conf()
+
+        # These 3 calls correspond to the first 3 mock_status calls above
+        self.assertFalse((yield self.service.is_running()))
+        self.assertFalse((yield self.service.is_running()))
+        self.assertTrue((yield self.service.is_running()))
+
+    @inlineCallbacks
+    def test_is_stable_yes(self):
+        self.setup_mock()
+        self.mock_check_success()
+        self.mocker.replay()
+
+        self.write_dummy_conf()
+        self.assertTrue((yield self.service.is_stable()))
+
+    @inlineCallbacks
+    def test_is_stable_no(self):
+        self.setup_mock()
+        self.mock_check_unstable()
+        self.mocker.replay()
+
+        self.write_dummy_conf()
+        self.assertFalse((yield self.service.is_stable()))
+
+    @inlineCallbacks
+    def test_is_stable_not_running(self):
+        self.setup_mock()
+        self.mock_check_not_running()
+        self.mocker.replay()
+
+        self.write_dummy_conf()
+        self.assertFalse((yield self.service.is_stable()))
+
+    @inlineCallbacks
+    def test_is_stable_not_even_installed(self):
+        self.assertFalse((yield self.service.is_stable()))
+
+    @inlineCallbacks
+    def test_get_pid(self):
+        """
+        Check get_pid interprets status output (when service is installed)
+        """
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_status(succeed("blah blob/gibbering blah"))
+        self.mock_status(succeed("blah start/running blah 12345"))
+        self.mocker.replay()
+
+        # Won't hit status; conf is not installed
+        self.assertEquals((yield self.service.get_pid()), None)
+        self.write_dummy_conf()
+
+        # These 3 calls correspond to the first 3 mock_status calls above
+        self.assertEquals((yield self.service.get_pid()), None)
+        self.assertEquals((yield self.service.get_pid()), None)
+        self.assertEquals((yield self.service.get_pid()), 12345)
+
+    @inlineCallbacks
+    def test_basic_install(self):
+        """Check a simple UpstartService writes expected conf file"""
+        e = yield self.assertFailure(self.service.install(), ServiceError)
+        self.assertEquals(str(e), "Cannot render .conf: no description set")
+        self.service.set_description("uninteresting service")
+        e = yield self.assertFailure(self.service.install(), ServiceError)
+        self.assertEquals(str(e), "Cannot render .conf: no command set")
+        self.service.set_command("/bin/false")
+        yield self.service.install()
+
+        self.assert_conf("test_basic_install")
+
+    @inlineCallbacks
+    def test_less_basic_install(self):
+        """Check conf for a different UpstartService (which sets an env var)"""
+        self.service.set_description("pew pew pew blam")
+        self.service.set_command("/bin/deathstar --ignore-ewoks endor")
+        self.service.set_environ({"FOO": "bar baz qux", "PEW": "pew"})
+        self.service.set_output_path("/somewhere/else")
+        yield self.service.install()
+
+        self.assert_conf("test_less_basic_install")
+
+    def test_install_via_script(self):
+        """Check that the output-as-script form does the right thing"""
+        self.setup_service()
+        install, start = self.service.get_cloud_init_commands()
+
+        os.system(install)
+        self.assert_conf()
+        self.assertEquals(start, "/sbin/start some-name")
+
+    @inlineCallbacks
+    def test_start_not_installed(self):
+        """Check that .start() also installs if necessary"""
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_start()
+        self.mock_check_success()
+        self.mocker.replay()
+
+        self.setup_service()
+        yield self.service.start()
+        self.assert_conf()
+
+    @inlineCallbacks
+    def test_start_not_started_stable(self):
+        """Check that .start() starts if stopped, and checks for stable pid"""
+        self.write_dummy_conf()
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_start("ignored")
+        self.mock_check_success()
+        self.mocker.replay()
+
+        self.setup_service()
+        yield self.service.start()
+        self.assert_dummy_conf()
+
+    @inlineCallbacks
+    def test_start_not_started_unstable(self):
+        """Check that .start() starts if stopped, and raises on unstable pid"""
+        self.write_dummy_conf()
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_start("kangaroo")
+        self.mock_check_unstable()
+        self.mocker.replay()
+
+        self.setup_service()
+        e = yield self.assertFailure(self.service.start(), ServiceError)
+        self.assertEquals(
+            str(e), "Failed to start job some-name; got output:\nkangaroo")
+        self.assert_dummy_conf()
+
+    @inlineCallbacks
+    def test_start_not_started_failure(self):
+        """Check that .start() starts if stopped, and raises on no pid"""
+        self.write_dummy_conf()
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_start()
+        self.mock_check_not_running()
+        self.mocker.replay()
+
+        self.setup_service()
+        e = yield self.assertFailure(self.service.start(), ServiceError)
+        self.assertEquals(
+            str(e), "Failed to start job some-name; no output detected")
+        self.assert_dummy_conf()
+
+    @inlineCallbacks
+    def test_start_started(self):
+        """Check that .start() does nothing if already running"""
+        self.write_dummy_conf()
+        self.setup_mock()
+        self.mock_status(succeed("blah start/running blah 12345"))
+        self.mocker.replay()
+
+        self.setup_service()
+        yield self.service.start()
+        self.assert_dummy_conf()
+
+    @inlineCallbacks
+    def test_destroy_not_installed(self):
+        """Check .destroy() does nothing if not installed"""
+        yield self.service.destroy()
+        self.assert_no_conf()
+
+    @inlineCallbacks
+    def test_destroy_not_started(self):
+        """Check .destroy just deletes conf if not running"""
+        self.write_dummy_conf()
+        self.setup_mock()
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mocker.replay()
+
+        yield self.service.destroy()
+        self.assert_no_conf()
+
+    @inlineCallbacks
+    def test_destroy_started(self):
+        """Check .destroy() stops running service and deletes conf file"""
+        self.write_dummy_conf()
+        self.setup_mock()
+        self.mock_status(succeed("blah start/running blah 54321"))
+        self.mock_stop()
+        self.mocker.replay()
+
+        yield self.service.destroy()
+        self.assert_no_conf()
+
+    @inlineCallbacks
+    def test_use_sudo(self):
+        """Check that expected commands are generated when use_sudo is set"""
+        self.setup_mock()
+        self.service = UpstartService("some-name", use_sudo=True)
+        self.setup_service()
+        with open(self.output, "w") as f:
+            f.write("clear this file out...")
+
+        def verify_cp(args, **kwargs):
+            sudo, cp, src, dst = args
+            self.assertEquals(sudo, "sudo")
+            self.assertEquals(cp, "cp")
+            with open(os.path.join(DATA_DIR, "test_standard_install")) as exp:
+                with open(src) as actual:
+                    self.assertEquals(actual.read(), exp.read())
+            self.assertEquals(dst, self.conf)
+            self.write_dummy_conf()
+
+        self.check_call(ANY, KWARGS)
+        self.mocker.call(verify_cp)
+        self.mock_call(("sudo", "rm", self.output))
+        self.mock_call(("sudo", "chmod", "644", self.conf))
+        self.mock_status(succeed("blah stop/waiting blah"))
+        self.mock_call(("sudo", "/sbin/start", "some-name"))
+        # 5 for initial stability check; 1 for final do-we-need-to-stop check
+        for _ in range(6):
+            self.mock_status(succeed("blah start/running blah 12345"))
+        self.mock_call(("sudo", "/sbin/stop", "some-name"))
+        self.mock_call(("sudo", "rm", self.conf))
+        self.mock_call(("sudo", "rm", self.output))
+
+        self.mocker.replay()
+        yield self.service.start()
+        yield self.service.destroy()

=== added file 'juju/lib/upstart.py'
--- juju/lib/upstart.py	1970-01-01 00:00:00 +0000
+++ juju/lib/upstart.py	2012-02-02 16:42:42 +0000
@@ -0,0 +1,166 @@
+import os
+import subprocess
+from tempfile import NamedTemporaryFile
+
+from twisted.internet.defer import inlineCallbacks, returnValue
+from twisted.internet.threads import deferToThread
+from twisted.internet.utils import getProcessOutput
+
+from juju.errors import ServiceError
+from juju.lib.twistutils import sleep
+
+
+_CONF_TEMPLATE = """\
+description "%s"
+author "Juju Team <[e-mail]>"
+
+start on runlevel [2345]
+stop on runlevel [!2345]
+respawn
+
+%s
+
+exec %s >> %s 2>&1
+"""
+
+def _silent_check_call(args):
+    with open(os.devnull, "w") as f:
+        return subprocess.check_call(
+            args, stdout=f.fileno(), stderr=f.fileno())
+
+
+class UpstartService(object):
+
+    # on class for ease of testing
+    init_dir = "/etc/init"
+
+    def __init__(self, name, init_dir=None, use_sudo=False):
+        self._name = name
+        if init_dir is not None:
+            self.init_dir = init_dir
+        self._use_sudo = use_sudo
+        self._output_path = None
+        self._description = None
+        self._environ = {}
+        self._command = None
+
+    @property
+    def _conf_path(self):
+        return os.path.join(
+            self.init_dir, "%s.conf" % self._name)
+
+    @property
+    def output_path(self):
+        if self._output_path is not None:
+            return self._output_path
+        return "/tmp/%s.output" % self._name
+
+    def set_description(self, description):
+        self._description = description
+
+    def set_environ(self, environ):
+        self._environ = environ
+
+    def set_command(self, command):
+        self._command = command
+
+    def set_output_path(self, path):
+        self._output_path = path
+
+    @inlineCallbacks
+    def _trash_output(self):
+        if os.path.exists(self.output_path):
+            # Just using os.unlink will fail when we're running TEST_SUDO tests
+            # which hit this code path (because root will own self.output_path)
+            yield self._call("rm", self.output_path)
+
+    def _render(self):
+        if self._description is None:
+            raise ServiceError("Cannot render .conf: no description set")
+        if self._command is None:
+            raise ServiceError("Cannot render .conf: no command set")
+        return _CONF_TEMPLATE % (
+            self._description,
+            "\n".join('env %s="%s"' % kv
+                      for kv in sorted(self._environ.items())),
+            self._command,
+            self.output_path)
+
+    def _call(self, *args):
+        if self._use_sudo:
+            args = ("sudo",) + args
+        return deferToThread(_silent_check_call, args)
+
+    def get_cloud_init_commands(self):
+        return ["cat >> %s <<EOF\n%sEOF\n" % (self._conf_path, self._render()),
+                "/sbin/start %s" % self._name]
+
+    @inlineCallbacks
+    def install(self):
+        with NamedTemporaryFile() as f:
+            f.write(self._render())
+            f.flush()
+            yield self._call("cp", f.name, self._conf_path)
+        yield self._call("chmod", "644", self._conf_path)
+
+    @inlineCallbacks
+    def start(self):
+        if not self.is_installed():
+            yield self.install()
+        if (yield self.is_running()):
+            return
+        yield self._trash_output()
+        yield self._call("/sbin/start", self._name)
+        if (yield self.is_stable()):
+            return
+
+        output = None
+        if os.path.exists(self.output_path):
+            with open(self.output_path) as f:
+                output = f.read()
+        if not output:
+            raise ServiceError(
+                "Failed to start job %s; no output detected" % self._name)
+        raise ServiceError(
+            "Failed to start job %s; got output:\n%s" % (self._name, output))
+
+    @inlineCallbacks
+    def destroy(self):
+        if (yield self.is_running()):
+            yield self._call("/sbin/stop", self._name)
+        if self.is_installed():
+            yield self._call("rm", self._conf_path)
+        yield self._trash_output()
+
+    @inlineCallbacks
+    def get_pid(self):
+        if not self.is_installed():
+            returnValue(None)
+        status = yield getProcessOutput("/sbin/status", [self._name])
+        if "start/running" not in status:
+            returnValue(None)
+        pid = status.split(" ")[-1]
+        returnValue(int(pid))
+
+    @inlineCallbacks
+    def is_running(self):
+        pid = yield self.get_pid()
+        returnValue(pid is not None)
+
+    @inlineCallbacks
+    def is_stable(self):
+        """Does the process continue to run with the same pid?
+
+        (5 times in a row, with a gap of 0.1s between each check)
+        """
+        pid = yield self.get_pid()
+        if pid is None:
+            returnValue(False)
+        for _ in range(4):
+            yield sleep(0.1)
+            if pid != (yield self.get_pid()):
+                returnValue(False)
+        returnValue(True)
+
+    def is_installed(self):
+        return os.path.exists(self._conf_path)

=== modified file 'juju/machine/tests/test_unit_deployment.py'
--- juju/machine/tests/test_unit_deployment.py	2011-10-05 13:59:44 +0000
+++ juju/machine/tests/test_unit_deployment.py	2012-02-02 16:42:42 +0000
@@ -3,25 +3,21 @@
 """
 import logging
 import os
-import sys
+import subprocess
 
-from twisted.internet.protocol import ProcessProtocol
 from twisted.internet.defer import inlineCallbacks, succeed
 
-import juju
 from juju.charm import get_charm_from_path
 from juju.charm.tests.test_repository import RepositoryTestBase
 from juju.lib.lxc import LXCContainer
-from juju.lib.mocker import MATCH, ANY
-from juju.lib.twistutils import get_module_directory
+from juju.lib.lxc.tests.test_lxc import uses_sudo
+from juju.lib.mocker import ANY, KWARGS
+from juju.lib.upstart import UpstartService
 from juju.machine.unit import UnitMachineDeployment, UnitContainerDeployment
 from juju.machine.errors import UnitDeploymentError
 from juju.tests.common import get_test_zookeeper_address
 
 
-MATCH_PROTOCOL = MATCH(lambda x: isinstance(x, ProcessProtocol))
-
-
 class UnitMachineDeploymentTest(RepositoryTestBase):
 
     def setUp(self):
@@ -32,6 +28,11 @@
         self.units_directory = os.path.join(self.juju_directory, "units")
         os.mkdir(self.units_directory)
         self.unit_name = "wordpress/0"
+        self.rootfs = self.makeDir()
+        self.init_dir = os.path.join(self.rootfs, "etc", "init")
+        os.makedirs(self.init_dir)
+        self.real_init_dir = self.patch(
+            UpstartService, "init_dir", self.init_dir)
 
         self.deployment = UnitMachineDeployment(
             self.unit_name,
@@ -41,11 +42,43 @@
             self.deployment.unit_agent_module, "juju.agents.unit")
         self.deployment.unit_agent_module = "juju.agents.dummy"
 
-    def process_kill(self, pid):
-        try:
-            os.kill(pid, 9)
-        except OSError:
-            pass
+    def setup_mock(self):
+        self.check_call = self.mocker.replace("subprocess.check_call")
+        self.getProcessOutput = self.mocker.replace(
+            "twisted.internet.utils.getProcessOutput")
+
+    def mock_is_running(self, running):
+        self.getProcessOutput("/sbin/status", ["juju-wordpress-0"])
+        if running:
+            self.mocker.result(succeed(
+                "juju-wordpress-0 start/running, process 12345"))
+        else:
+            self.mocker.result(succeed("juju-wordpress-0 stop/waiting"))
+
+    def _without_sudo(self, args, **_):
+        self.assertEquals(args[0], "sudo")
+        return subprocess.call(args[1:])
+
+    def mock_install(self):
+        self.check_call(ANY, KWARGS)  # cp to init dir
+        self.mocker.call(self._without_sudo)
+        self.check_call(ANY, KWARGS)  # chmod 644
+        self.mocker.call(self._without_sudo)
+
+    def mock_start(self):
+        self.check_call(("sudo", "/sbin/start", "juju-wordpress-0"), KWARGS)
+        self.mocker.result(0)
+        for _ in range(5):
+            self.mock_is_running(True)
+
+    def mock_destroy(self):
+        self.check_call(("sudo", "/sbin/stop", "juju-wordpress-0"), KWARGS)
+        self.mocker.result(0)
+        self.check_call(ANY, KWARGS)  # rm from init dir
+        self.mocker.call(self._without_sudo)
+
+    def assert_pid_running(self, pid, expect):
+        self.assertEquals(os.path.exists("/proc/%s" % pid), expect)
 
     def test_unit_name_with_path_manipulation_raises_assertion(self):
         self.assertRaises(
@@ -60,129 +93,50 @@
             os.path.join(self.units_directory,
                          self.unit_name.replace("/", "-")))
 
-    def test_unit_pid_file(self):
-        self.assertEqual(
-            self.deployment.pid_file,
-            os.path.join(self.units_directory,
-                         "%s.pid" % (self.unit_name.replace("/", "-"))))
-
     def test_service_unit_start(self):
         """
         Starting a service unit will result in a unit workspace being created
         if it does not exist and a running service unit agent.
         """
-
-        d = self.deployment.start(
-            "0", get_test_zookeeper_address(), self.bundle)
-
-        @inlineCallbacks
-        def validate_result(result):
-            # give process time to write its pid
-            yield self.sleep(0.1)
-            self.addCleanup(
-                self.process_kill,
-                int(open(self.deployment.pid_file).read()))
-            self.assertEqual(result, True)
-
-        d.addCallback(validate_result)
-        return d
-
-    def test_deployment_get_environment(self):
-        zk_address = get_test_zookeeper_address()
-        environ = self.deployment.get_environment(21, zk_address)
-        environ.pop("PYTHONPATH")
-        self.assertEqual(environ["JUJU_HOME"], self.juju_directory)
-        self.assertEqual(environ["JUJU_UNIT_NAME"], self.unit_name)
-        self.assertEqual(environ["JUJU_ZOOKEEPER"], zk_address)
-        self.assertEqual(environ["JUJU_MACHINE_ID"], "21")
-
-    def test_service_unit_start_with_integer_machine_id(self):
-        """
-        Starting a service unit will result in a unit workspace being created
-        if it does not exist and a running service unit agent.
-        """
-        d = self.deployment.start(
-            21, get_test_zookeeper_address(), self.bundle)
-
-        @inlineCallbacks
-        def validate_result(result):
-            # give process time to write its pid
-            yield self.sleep(0.1)
-            self.addCleanup(
-                self.process_kill,
-                int(open(self.deployment.pid_file).read()))
-            self.assertEqual(result, True)
-
-        d.addCallback(validate_result)
-        return d
-
-    def test_service_unit_start_with_agent_startup_error(self):
-        """
-        Starting a service unit will result in a unit workspace being created
-        if it does not exist and a running service unit agent.
-        """
-        self.deployment.unit_agent_module = "magichat.xr1"
-        d = self.deployment.start(
-            "0", get_test_zookeeper_address(), self.bundle)
-
-        self.failUnlessFailure(d, UnitDeploymentError)
-
-        def validate_result(error):
-            self.assertIn("No module named magichat", str(error))
-
-        d.addCallback(validate_result)
-        return d
-
-    def test_service_unit_start_agent_arguments(self):
-        """
-        Starting a service unit will start a service unit agent with arguments
-        denoting the current machine id, zookeeper server location, and the
-        unit name. Additionally it will configure the log and pid file
-        locations.
-        """
-        machine_id = "0"
-        zookeeper_hosts = "menagerie.example.com:2181"
-
-        from twisted.internet import reactor
-        mock_reactor = self.mocker.patch(reactor)
-
-        environ = dict(os.environ)
-        environ["JUJU_UNIT_NAME"] = self.unit_name
-        environ["JUJU_HOME"] = self.juju_directory
-        environ["JUJU_MACHINE_ID"] = machine_id
-        environ["JUJU_ZOOKEEPER"] = zookeeper_hosts
-        environ["PYTHONPATH"] = ":".join(
-            filter(None, [
-                os.path.dirname(get_module_directory(juju)),
-                environ.get("PYTHONPATH")]))
-
-        pid_file = os.path.join(
-            self.units_directory,
-            "%s.pid" % self.unit_name.replace("/", "-"))
-
-        log_file = os.path.join(
-            self.deployment.directory,
-            "charm.log")
-
-        args = [sys.executable, "-m", "juju.agents.dummy", "-n",
-                "--pidfile", pid_file, "--logfile", log_file]
-
-        mock_reactor.spawnProcess(
-            MATCH_PROTOCOL, sys.executable, args, environ)
+        self.setup_mock()
+        self.mock_install()
+        self.mock_is_running(False)
+        self.mock_start()
         self.mocker.replay()
-        self.deployment.start(
-            machine_id, zookeeper_hosts, self.bundle)
-
-    def xtest_service_unit_start_pre_unpack(self):
-        """
-        Attempting to start a charm before the charm is unpacked
-        results in an exception.
-        """
-        error = yield self.assertFailure(
-            self.deployment.start(
-                "0", get_test_zookeeper_address(), self.bundle),
-            UnitDeploymentError)
-        self.assertEquals(str(error), "Charm must be unpacked first.")
+
+        d = self.deployment.start(
+            "123", get_test_zookeeper_address(), self.bundle)
+
+        def verify_upstart(_):
+            conf_path = os.path.join(self.init_dir, "juju-wordpress-0.conf")
+            with open(conf_path) as f:
+                lines = f.readlines()
+
+            env = []
+            for line in lines:
+                if line.startswith("env "):
+                    env.append(line[4:-1].split("=", 1))
+                if line.startswith("exec "):
+                    exec_ = line[5:-1]
+
+            env = dict((k, v.strip('"')) for (k, v) in env)
+            env.pop("PYTHONPATH")
+            self.assertEquals(env, {
+                "JUJU_HOME": self.juju_directory,
+                "JUJU_UNIT_NAME": self.unit_name,
+                "JUJU_ZOOKEEPER": get_test_zookeeper_address(),
+                "JUJU_MACHINE_ID": "123"})
+
+            log_file = os.path.join(
+                self.deployment.directory, "charm.log")
+            command = " ".join([
+                "/usr/bin/python", "-m", "juju.agents.dummy", "--nodaemon",
+                "--logfile", log_file, "--session-file",
+                "/var/run/juju/unit-wordpress-0-agent.zksession",
+                ">> /tmp/juju-wordpress-0.output 2>&1"])
+            self.assertEquals(exec_, command)
+        d.addCallback(verify_upstart)
+        return d
 
     @inlineCallbacks
     def test_service_unit_destroy(self):
@@ -190,48 +144,22 @@
         Forcibly stop a unit, and destroy any directories associated to it
         on the machine, and kills the unit agent process.
         """
+        self.setup_mock()
+        self.mock_install()
+        self.mock_is_running(False)
+        self.mock_start()
+        self.mock_is_running(True)
+        self.mock_destroy()
+        self.mocker.replay()
+
         yield self.deployment.start(
             "0", get_test_zookeeper_address(), self.bundle)
-        # give the process time to write its pid file
-        yield self.sleep(0.1)
-        pid = int(open(self.deployment.pid_file).read())
+
         yield self.deployment.destroy()
-        # give the process time to die.
-        yield self.sleep(0.1)
-        e = self.assertRaises(OSError, os.kill, pid, 0)
-        self.assertEqual(e.errno, 3)
         self.assertFalse(os.path.exists(self.deployment.directory))
-        self.assertFalse(os.path.exists(self.deployment.pid_file))
-
-    def test_service_unit_destroy_stale_pid(self):
-        """
-        A stale pid file does not cause any errors.
-
-        We mock away is_running as otherwise it will check for this, but
-        there exists a small window when the result may disagree.
-        """
-        self.makeFile("8917238", path=self.deployment.pid_file)
-        mock_deployment = self.mocker.patch(self.deployment)
-        mock_deployment.is_running()
-        self.mocker.result(succeed(True))
-        self.mocker.replay()
-        return self.deployment.destroy()
-
-    def test_service_unit_destroy_perm_error(self):
-        """
-        A stale pid file does not cause any errors.
-
-        We mock away is_running as otherwise it will check for this, but
-        there exists a small window when the result may disagree.
-        """
-        if os.geteuid() == 0:
-            return
-        self.makeFile("1", path=self.deployment.pid_file)
-        mock_deployment = self.mocker.patch(self.deployment)
-        mock_deployment.is_running()
-        self.mocker.result(succeed(True))
-        self.mocker.replay()
-        return self.assertFailure(self.deployment.destroy(), OSError)
+
+        conf_path = os.path.join(self.init_dir, "juju-wordpress-0.conf")
+        self.assertFalse(os.path.exists(conf_path))
 
     @inlineCallbacks
     def test_service_unit_destroy_undeployed(self):
@@ -247,11 +175,23 @@
         If the unit is not running, then destroy will just remove
         its directory.
         """
-        self.deployment.unpack_charm(self.bundle)
-        self.assertTrue(os.path.exists(self.deployment.directory))
+        self.setup_mock()
+        self.mock_install()
+        self.mock_is_running(False)
+        self.mock_start()
+        self.mock_is_running(False)
+        self.check_call(ANY, KWARGS)  # rm from init dir
+        self.mocker.call(self._without_sudo)
+        self.mocker.replay()
+
+        yield self.deployment.start(
+            "0", get_test_zookeeper_address(), self.bundle)
         yield self.deployment.destroy()
         self.assertFalse(os.path.exists(self.deployment.directory))
 
+        conf_path = os.path.join(self.init_dir, "juju-wordpress-0.conf")
+        self.assertFalse(os.path.exists(conf_path))
+
     def test_unpack_charm(self):
         """
         The deployment unpacks a charm bundle into the unit workspace.
@@ -279,62 +219,81 @@
             str(error),
             "Invalid charm for deployment: %s" % self.charm.path)
 
-    def test_is_running_no_pid_file(self):
-        """
-        If there is no pid file the service unit is not running.
-        """
-        self.assertEqual((yield self.deployment.is_running()), False)
-
-    def test_is_running(self):
-        """
-        The service deployment will check the pid and validate
-        that the pid found is a running process.
-        """
-        self.makeFile(
-            str(os.getpid()), path=self.deployment.pid_file)
+    @inlineCallbacks
+    def test_is_running_not_installed(self):
+        """
+        If there is no conf file the service unit is not running.
+        """
+        self.assertEqual((yield self.deployment.is_running()), False)
+
+    @inlineCallbacks
+    def test_is_running_not_running(self):
+        """
+        If the conf file exists, but job not running, unit not running
+        """
+        conf_path = os.path.join(self.init_dir, "juju-wordpress-0.conf")
+        with open(conf_path, "w") as f:
+            f.write("blah")
+        self.setup_mock()
+        self.mock_is_running(False)
+        self.mocker.replay()
+        self.assertEqual((yield self.deployment.is_running()), False)
+
+    @inlineCallbacks
+    def test_is_running_success(self):
+        """
+        Check running job.
+        """
+        conf_path = os.path.join(self.init_dir, "juju-wordpress-0.conf")
+        with open(conf_path, "w") as f:
+            f.write("blah")
+        self.setup_mock()
+        self.mock_is_running(True)
+        self.mocker.replay()
         self.assertEqual((yield self.deployment.is_running()), True)
 
-    def test_is_running_against_unknown_error(self):
-        """
-        If we don't have permission to access the process, the
-        original error should get passed along.
-        """
-        if os.geteuid() == 0:
-            return
-        self.makeFile("1", path=self.deployment.pid_file)
-        self.assertFailure(self.deployment.is_running(), OSError)
-
-    def test_is_running_invalid_pid_file(self):
-        """
-        If the pid file is corrupted on disk, and does not contain
-        a valid integer, then the agent is not running.
-        """
-        self.makeFile("abcdef", path=self.deployment.pid_file)
-        self.assertEqual(
-            (yield self.deployment.is_running()), False)
-
-    def test_is_running_invalid_pid(self):
-        """
-        If the pid file refers to an invalid process then the
-        agent is not running.
-        """
-        self.makeFile("669966", path=self.deployment.pid_file)
-        self.assertEqual(
-            (yield self.deployment.is_running()), False)
-
-
-upstart_job_sample = '''\
-description "Unit agent for riak/0"
-author "Juju Team <[e-mail]>"
-start on start on filesystem or runlevel [2345]
-stop on runlevel [!2345]
-
-respawn
-
-env JUJU_MACHINE_ID="0"
-env JUJU_HOME="/var/lib/juju"
-env JUJU_ZOOKEEPER="127.0.1.1:2181"
-env JUJU_UNIT_NAME="riak/0"'''
+    @uses_sudo
+    @inlineCallbacks
+    def test_run_actual_process(self):
+        # "unpatch" to use real /etc/init
+        self.patch(UpstartService, "init_dir", self.real_init_dir)
+        yield self.deployment.start(
+            "0", get_test_zookeeper_address(), self.bundle)
+        old_pid = yield self.deployment.get_pid()
+        self.assert_pid_running(old_pid, True)
+
+        # Give the job a chance to fall over and be restarted (if the
+        # pid doesn't change, that hasn't hapened)
+        yield self.sleep(0.1)
+        self.assertEquals((yield self.deployment.get_pid()), old_pid)
+        self.assert_pid_running(old_pid, True)
+
+        # Kick the job over ourselves; check it comes back
+        os.system("sudo kill -9 %s" % old_pid)
+        yield self.sleep(0.1)
+        self.assert_pid_running(old_pid, False)
+        new_pid = yield self.deployment.get_pid()
+        self.assertNotEquals(new_pid, old_pid)
+        self.assert_pid_running(new_pid, True)
+
+        yield self.deployment.destroy()
+        self.assertEquals((yield self.deployment.get_pid()), None)
+        self.assert_pid_running(new_pid, False)
+
+    @uses_sudo
+    @inlineCallbacks
+    def test_fail_to_run_actual_process(self):
+        self.deployment.unit_agent_module = "haha.disregard.that"
+        self.patch(UpstartService, "init_dir", self.real_init_dir)
+
+        d = self.deployment.start(
+            "0", get_test_zookeeper_address(), self.bundle)
+        e = yield self.assertFailure(d, UnitDeploymentError)
+        self.assertTrue(str(e).startswith(
+            "Failed to start job juju-wordpress-0; got output:\n"))
+        self.assertIn("No module named haha", str(e))
+
+        yield self.deployment.destroy()
 
 
 class UnitContainerDeploymentTest(RepositoryTestBase):
@@ -365,14 +324,6 @@
             "ns1-riak-0",
             self.unit_deploy.container_name)
 
-    def test_get_upstart_job(self):
-        upstart_job = self.unit_deploy.get_upstart_unit_job(
-            0, "127.0.1.1:2181")
-        job = self.get_normalized(upstart_job)
-        self.assertIn('JUJU_ZOOKEEPER="127.0.1.1:2181"', job)
-        self.assertIn('JUJU_MACHINE_ID="0"', job)
-        self.assertIn('JUJU_UNIT_NAME="riak/0"', job)
-
     @inlineCallbacks
     def test_destroy(self):
         mock_container = self.mocker.patch(self.unit_deploy.container)
@@ -403,7 +354,7 @@
         unit_deploy = UnitContainerDeployment(
             self.unit_name, self.juju_home)
         container = yield unit_deploy._get_master_template(
-            "local", "127.0.0.1:1", "abc")
+            "local", "abc")
         self.assertEqual(container.origin, "lp:~juju/foobar")
         self.assertEqual(
             container.customize_log,
@@ -420,7 +371,7 @@
         mock_deploy = self.mocker.patch(self.unit_deploy)
         # this minimally validates that we are also called with the
         # expect public key
-        mock_deploy._get_container(ANY, ANY, ANY, env["JUJU_PUBLIC_KEY"])
+        mock_deploy._get_container(ANY, ANY, env["JUJU_PUBLIC_KEY"])
         self.mocker.result((container, rootfs))
 
         mock_container = self.mocker.patch(container)
@@ -434,7 +385,7 @@
         yield self.unit_deploy.start("0", "127.0.1.1:2181", self.bundle)
 
         # Verify the upstart job
-        upstart_agent_name = "%s-unit-agent.conf" % (
+        upstart_agent_name = "juju-%s.conf" % (
             self.unit_name.replace("/", "-"))
         content = open(
             os.path.join(rootfs, "etc", "init", upstart_agent_name)).read()
@@ -443,10 +394,13 @@
         self.assertIn('JUJU_MACHINE_ID="0"', job)
         self.assertIn('JUJU_UNIT_NAME="riak/0"', job)
 
-        # Verify the symlink exists
+        # Verify the symlinks exist
         self.assertTrue(os.path.lexists(os.path.join(
             self.unit_deploy.juju_home, "units",
             self.unit_deploy.unit_path_name, "unit.log")))
+        self.assertTrue(os.path.lexists(os.path.join(
+            self.unit_deploy.juju_home, "units",
+            self.unit_deploy.unit_path_name, "output.log")))
 
         # Verify the charm is on disk.
         self.assertTrue(os.path.exists(os.path.join(
@@ -471,7 +425,7 @@
         container = LXCContainer(self.unit_name, None, None, None)
 
         mock_deploy = self.mocker.patch(self.unit_deploy)
-        mock_deploy._get_master_template(ANY, ANY, ANY)
+        mock_deploy._get_master_template(ANY, ANY)
         self.mocker.result(container)
 
         mock_container = self.mocker.patch(container)
@@ -481,7 +435,7 @@
         self.mocker.replay()
 
         container, rootfs = yield self.unit_deploy._get_container(
-            "0", "127.0.0.1:2181", None, "dsa...")
+            "0", None, "dsa...")
 
         output = self.output.getvalue()
         self.assertIn("Container created for %s" % self.unit_deploy.unit_name,

=== modified file 'juju/machine/unit.py'
--- juju/machine/unit.py	2011-10-01 00:04:14 +0000
+++ juju/machine/unit.py	2012-02-02 16:42:42 +0000
@@ -1,19 +1,17 @@
 import os
-import errno
-import signal
 import shutil
 import sys
 import logging
 
 import juju
 
-from twisted.internet.defer import (
-    Deferred, inlineCallbacks, returnValue, succeed, fail)
-from twisted.internet.protocol import ProcessProtocol
+from twisted.internet.defer import inlineCallbacks, returnValue
 
 from juju.charm.bundle import CharmBundle
+from juju.errors import ServiceError
+from juju.lib.lxc import LXCContainer, get_containers, LXCError
 from juju.lib.twistutils import get_module_directory
-from juju.lib.lxc import LXCContainer, get_containers, LXCError
+from juju.lib.upstart import UpstartService
 
 from .errors import UnitDeploymentError
 
@@ -26,22 +24,17 @@
     return UnitMachineDeployment
 
 
-class AgentProcessProtocol(ProcessProtocol):
-
-    def __init__(self, deferred):
-        self.deferred = deferred
-        self._error_buffer = []
-
-    def errReceived(self, data):
-        self._error_buffer.append(data)
-
-    def processEnded(self, reason):
-        if self._error_buffer:
-            msg = "".join(self._error_buffer)
-            msg.strip()
-            self.deferred.errback(UnitDeploymentError(msg))
-        else:
-            self.deferred.callback(True)
+def _get_environment(unit_name, juju_home, machine_id, zookeeper_hosts):
+    environ = dict()
+    environ["JUJU_MACHINE_ID"] = str(machine_id)
+    environ["JUJU_UNIT_NAME"] = unit_name
+    environ["JUJU_HOME"] = juju_home
+    environ["JUJU_ZOOKEEPER"] = zookeeper_hosts
+    environ["PYTHONPATH"] = ":".join(
+        filter(None, [
+            os.path.dirname(get_module_directory(juju)),
+            os.environ.get("PYTHONPATH")]))
+    return environ
 
 
 class UnitMachineDeployment(object):
@@ -58,46 +51,35 @@
     unit_agent_module = "juju.agents.unit"
 
     def __init__(self, unit_name, juju_home):
-        self.unit_name = unit_name
-
         assert ".." not in unit_name, "Invalid Unit Name"
+        self.unit_name = unit_name
+        self.juju_home = juju_home
         self.unit_path_name = unit_name.replace("/", "-")
-        self.juju_home = juju_home
-
         self.directory = os.path.join(
             self.juju_home, "units", self.unit_path_name)
-
-        self.pid_file = os.path.join(
-            self.juju_home, "units", "%s.pid" % self.unit_path_name)
-
+        self.service = UpstartService(
+            # NOTE: we need use_sudo to work correctly during tests that
+            # launch actual processes (rather than just mocking/trusting).
+            "juju-%s" % self.unit_path_name, use_sudo=True)
+
+    @inlineCallbacks
     def start(self, machine_id, zookeeper_hosts, bundle):
         """Start a service unit agent."""
-        # Extract the charm into the unit directory.
         self.unpack_charm(bundle)
-
-        # Start the service unit agent
-        log_file = os.path.join(self.directory, "charm.log")
-        environ = self.get_environment(machine_id, zookeeper_hosts)
-        args = [sys.executable, "-m", self.unit_agent_module, "-n",
-                "--pidfile", self.pid_file, "--logfile", log_file]
-
-        from twisted.internet import reactor
-        process_deferred = Deferred()
-        protocol = AgentProcessProtocol(process_deferred)
-        reactor.spawnProcess(protocol, sys.executable, args, environ)
-        return process_deferred
-
-    def get_environment(self, machine_id, zookeeper_hosts):
-        environ = dict(os.environ)
-        environ["JUJU_MACHINE_ID"] = str(machine_id)
-        environ["JUJU_UNIT_NAME"] = self.unit_name
-        environ["JUJU_HOME"] = self.juju_home
-        environ["JUJU_ZOOKEEPER"] = zookeeper_hosts
-        environ["PYTHONPATH"] = ":".join(
-            filter(None, [
-                os.path.dirname(get_module_directory(juju)),
-                environ.get("PYTHONPATH")]))
-        return environ
+        self.service.set_description(
+            "Juju unit agent for %s" % self.unit_name)
+        self.service.set_environ(_get_environment(
+            self.unit_name, self.juju_home, machine_id, zookeeper_hosts))
+        self.service.set_command(" ".join((
+            "/usr/bin/python", "-m", self.unit_agent_module,
+            "--nodaemon",
+            "--logfile", os.path.join(self.directory, "charm.log"),
+            "--session-file",
+            "/var/run/juju/unit-%s-agent.zksession" % self.unit_path_name)))
+        try:
+            yield self.service.start()
+        except ServiceError as e:
+            raise UnitDeploymentError(str(e))
 
     @inlineCallbacks
     def destroy(self):
@@ -105,41 +87,17 @@
 
         This will destroy/unmount any state on disk.
         """
-        running = yield self.is_running()
-        if running:
-            pid = int(open(self.pid_file).read())
-            try:
-                os.kill(pid, signal.SIGKILL)
-            except OSError, e:
-                if e.errno != errno.ESRCH:
-                    raise
-
-        if os.path.exists(self.pid_file):
-            os.remove(self.pid_file)
-
+        yield self.service.destroy()
         if os.path.exists(self.directory):
             shutil.rmtree(self.directory)
 
+    def get_pid(self):
+        """Get the service unit's process id."""
+        return self.service.get_pid()
+
     def is_running(self):
         """Is the service unit running."""
-        try:
-            with open(self.pid_file) as pid_fh:
-                pid = int(pid_fh.read())
-        except (IOError, ValueError):
-            return succeed(False)
-
-        # Attempt to send a signal to the process to verify its a valid process
-        # From man 2 kill
-        # "If sig is 0, then no signal is sent, but error checking is still
-        # performed; this can be used to check for the existence of a process
-        # ID or process group ID."
-        try:
-            os.kill(pid, 0)
-        except OSError, e:
-            if e.errno == errno.ESRCH:
-                return succeed(False)
-            return fail(e)
-        return succeed(True)
+        return self.service.is_running()
 
     def unpack_charm(self, charm):
         """Unpack a charm to the service units directory."""
@@ -150,27 +108,7 @@
         charm.extract_to(os.path.join(self.directory, "charm"))
 
 
-container_upstart_job_template = """\
-description "Unit agent for %(JUJU_UNIT_NAME)s"
-author "Juju Team <[e-mail]>"
-
-start on start on filesystem or runlevel [2345]
-stop on runlevel [!2345]
-
-respawn
-
-env JUJU_MACHINE_ID="%(JUJU_MACHINE_ID)s"
-env JUJU_HOME="%(JUJU_HOME)s"
-env JUJU_ZOOKEEPER="%(JUJU_ZOOKEEPER)s"
-env JUJU_UNIT_NAME="%(JUJU_UNIT_NAME)s"
-env PYTHONPATH="%(PYTHONPATH)s"
-
-exec /usr/bin/python -m juju.agents.unit \
-     --logfile=/var/log/juju/unit-%(UNIT_PATH_NAME)s.log
-"""
-
-
-class UnitContainerDeployment(UnitMachineDeployment):
+class UnitContainerDeployment(object):
     """Deploy a service unit in a container.
 
     Units deployed in a container have strong isolation between
@@ -185,66 +123,35 @@
     """
 
     def __init__(self, unit_name, juju_home):
-        super(UnitContainerDeployment, self).__init__(unit_name, juju_home)
+        self.unit_name = unit_name
+        self.juju_home = juju_home
+        self.unit_path_name = unit_name.replace("/", "-")
 
+        self._juju_origin = os.environ.get("JUJU_ORIGIN")
         self._unit_namespace = os.environ.get("JUJU_UNIT_NS")
-        self._juju_origin = os.environ.get("JUJU_ORIGIN")
         assert self._unit_namespace is not None, "Required unit ns not found"
+        self.container_name = "%s-%s" % (
+            self._unit_namespace, self.unit_path_name)
 
-        self.pid_file = None
         self.container = LXCContainer(self.container_name, None, None, None)
-
-    @property
-    def container_name(self):
-        """Get a qualfied name for the container.
-
-        The units directory for the machine points to a path like::
-
-           /var/lib/juju/units
-
-        In the case of the local provider this directory is qualified
-        to allow for multiple users with multiple environments::
-
-          /var/lib/juju/username-envname
-
-        This value is passed to the agent via the JUJU_HOME environment
-        variable.
-
-        This function extracts the name qualifier for the container from
-        the JUJU_HOME value.
-        """
-        return "%s-%s" % (self._unit_namespace,
-                          self.unit_name.replace("/", "-"))
+        self.directory = None
 
     def setup_directories(self):
         # Create state directories for unit in the container
         # Move to juju-create script
-        units_dir = os.path.join(
-            self.directory, "var", "lib", "juju", "units")
-        if not os.path.exists(units_dir):
-            os.makedirs(units_dir)
-
-        state_dir = os.path.join(
-            self.directory, "var", "lib", "juju", "state")
-        if not os.path.exists(state_dir):
-            os.makedirs(state_dir)
-
-        log_dir = os.path.join(
-            self.directory, "var", "log", "juju")
-        if not os.path.exists(log_dir):
-            os.makedirs(log_dir)
-
-        unit_dir = os.path.join(units_dir, self.unit_path_name)
-        if not os.path.exists(unit_dir):
-            os.mkdir(unit_dir)
-
-        host_unit_dir = os.path.join(
-            self.juju_home, "units", self.unit_path_name)
-        if not os.path.exists(host_unit_dir):
-            os.makedirs(host_unit_dir)
+        base = self.directory
+        dirs = ((base, "var", "lib", "juju", "units", self.unit_path_name),
+                (base, "var", "lib", "juju", "state"),
+                (base, "var", "log", "juju"),
+                (self.juju_home, "units", self.unit_path_name))
+
+        for parts in dirs:
+            dir_ = os.path.join(*parts)
+            if not os.path.exists(dir_):
+                os.makedirs(dir_)
 
     @inlineCallbacks
-    def _get_master_template(self, machine_id, zookeeper_hosts, public_key):
+    def _get_master_template(self, machine_id, public_key):
         container_template_name = "%s-%s-template" % (
             self._unit_namespace, machine_id)
 
@@ -260,7 +167,7 @@
         if not master_template.is_constructed():
             log.debug("Creating master container...")
             yield master_template.create()
-            log.debug("Created master container %s" % container_template_name)
+            log.debug("Created master container %s", container_template_name)
 
         # it wasn't constructed and we couldn't construct it
         if not master_template.is_constructed():
@@ -269,15 +176,15 @@
         returnValue(master_template)
 
     @inlineCallbacks
-    def _get_container(self, machine_id, zookeeper_hosts, bundle, public_key):
+    def _get_container(self, machine_id, bundle, public_key):
         master_template = yield self._get_master_template(
-            machine_id, zookeeper_hosts, public_key)
+            machine_id, public_key)
         log.info(
-            "Creating container %s...", os.path.basename(self.directory))
+            "Creating container %s...", self.unit_path_name)
 
         container = yield master_template.clone(self.container_name)
         directory = container.rootfs
-        log.info("Container created for %s" % self.unit_name)
+        log.info("Container created for %s", self.unit_name)
         returnValue((container, directory))
 
     @inlineCallbacks
@@ -293,10 +200,9 @@
         # Build a template container that can be cloned in deploy
         # we leave the loosely initialized self.container in place for
         # the class as thats all we need for methods other than start.
-        self.container, self.directory = yield self._get_container(machine_id,
-                                                         zookeeper_hosts,
-                                                         bundle,
-                                                         public_key)
+        self.container, self.directory = yield self._get_container(
+            machine_id, bundle, public_key)
+
         # Create state directories for unit in the container
         self.setup_directories()
 
@@ -308,13 +214,25 @@
         log.debug("Charm extracted into container")
 
         # Write upstart file for the agent into the container
-        upstart_path = os.path.join(
-            self.directory, "etc", "init",
-            "%s-unit-agent.conf" % self.unit_path_name)
-        with open(upstart_path, "w") as fh:
-            fh.write(self.get_upstart_unit_job(machine_id, zookeeper_hosts))
+        service_name = "juju-%s" % self.unit_path_name
+        init_dir = os.path.join(self.directory, "etc", "init")
+        service = UpstartService(service_name, init_dir=init_dir)
+        service.set_description(
+            "Juju unit agent for %s" % self.unit_name)
+        service.set_environ(_get_environment(
+            self.unit_name, "/var/lib/juju", machine_id, zookeeper_hosts))
+        service.set_output_path(
+            "/var/log/juju/unit-%s-output.log" % self.unit_path_name)
+        service.set_command(" ".join((
+            "/usr/bin/python",
+            "-m", "juju.agents.unit",
+            "--nodaemon",
+            "--logfile", "/var/log/juju/unit-%s.log" % self.unit_path_name,
+            "--session-file",
+            "/var/run/juju/unit-%s-agent.zksession" % self.unit_path_name)))
+        yield service.install()
 
-        # Create a symlink on the host for easier access to the unit log file
+        # Create symlinks on the host for easier access to the unit log files
         unit_log_path_host = os.path.join(
             self.juju_home, "units", self.unit_path_name, "unit.log")
         if not os.path.lexists(unit_log_path_host):
@@ -322,6 +240,13 @@
                 os.path.join(self.directory, "var", "log", "juju",
                              "unit-%s.log" % self.unit_path_name),
                 unit_log_path_host)
+        unit_output_path_host = os.path.join(
+            self.juju_home, "units", self.unit_path_name, "output.log")
+        if not os.path.lexists(unit_output_path_host):
+            os.symlink(
+                os.path.join(self.directory, "var", "log", "juju",
+                             "unit-%s-output.log" % self.unit_path_name),
+                unit_output_path_host)
 
         # Debug log for the container
         container_log_path = os.path.join(
@@ -330,36 +255,22 @@
 
         log.debug("Starting container...")
         yield self.container.run()
-        log.info("Started container for %s" % self.unit_name)
+        log.info("Started container for %s", self.unit_name)
 
     @inlineCallbacks
     def destroy(self):
-        """Destroy the unit container.
-        """
+        """Destroy the unit container."""
         log.debug("Destroying container...")
         yield self.container.destroy()
-        log.info("Destroyed container for %s" % self.unit_name)
+        log.info("Destroyed container for %s", self.unit_name)
 
     @inlineCallbacks
     def is_running(self):
-        """Is the unit container running.
-        """
-        # TODO: container running may not imply agent running.  the
-        # pid file has the pid from the container, we need a container
-        # pid -> host pid mapping to query status from the machine agent.
-        # alternatively querying zookeeper for the unit agent presence
-        # node.
+        """Is the unit container running?"""
+        # TODO: container running may not imply agent running.
+        # query zookeeper for the unit agent presence node?
         if not self.container:
             returnValue(False)
         container_map = yield get_containers(
             prefix=self.container.container_name)
         returnValue(container_map.get(self.container.container_name, False))
-
-    def get_upstart_unit_job(self, machine_id, zookeeper_hosts):
-        """Return a string containing the  upstart job to start the unit agent.
-        """
-        environ = self.get_environment(machine_id, zookeeper_hosts)
-        # Keep qualified locations within the container for colo support
-        environ["JUJU_HOME"] = "/var/lib/juju"
-        environ["UNIT_PATH_NAME"] = self.unit_path_name
-        return container_upstart_job_template % environ

=== modified file 'juju/providers/common/cloudinit.py'
--- juju/providers/common/cloudinit.py	2012-01-09 13:58:21 +0000
+++ juju/providers/common/cloudinit.py	2012-02-02 16:42:42 +0000
@@ -1,6 +1,7 @@
 from subprocess import Popen, PIPE
 
 from juju.errors import CloudInitError
+from juju.lib.upstart import UpstartService
 from juju.providers.common.utils import format_cloud_init
 from juju.state.auth import make_identity
 import juju
@@ -41,21 +42,26 @@
 
 
 def _machine_scripts(machine_id, zookeeper_hosts):
-    return [
-        "JUJU_MACHINE_ID=%s JUJU_ZOOKEEPER=%s "
-        "python -m juju.agents.machine -n "
-        "--logfile=/var/log/juju/machine-agent.log "
-        "--pidfile=/var/run/juju/machine-agent.pid"
-        % (machine_id, zookeeper_hosts)]
+    service = UpstartService("juju-machine-agent")
+    service.set_description("Juju machine agent")
+    service.set_environ(
+        {"JUJU_MACHINE_ID": machine_id, "JUJU_ZOOKEEPER": zookeeper_hosts})
+    service.set_command(
+        "python -m juju.agents.machine --nodaemon "
+        "--logfile /var/log/juju/machine-agent.log "
+        "--session-file /var/run/juju/machine-agent.zksession")
+    return service.get_cloud_init_commands()
 
 
 def _provision_scripts(zookeeper_hosts):
-    return [
-        "JUJU_ZOOKEEPER=%s "
-        "python -m juju.agents.provision -n "
-        "--logfile=/var/log/juju/provision-agent.log "
-        "--pidfile=/var/run/juju/provision-agent.pid"
-        % zookeeper_hosts]
+    service = UpstartService("juju-provision-agent")
+    service.set_description("Juju provisioning agent")
+    service.set_environ({"JUJU_ZOOKEEPER": zookeeper_hosts})
+    service.set_command(
+        "python -m juju.agents.provision --nodaemon "
+        "--logfile /var/log/juju/provision-agent.log "
+        "--session-file /var/run/juju/provision-agent.zksession")
+    return service.get_cloud_init_commands()
 
 
 def _line_generator(data):
@@ -64,6 +70,7 @@
         if stripped:
             yield (len(line)-len(stripped), stripped)
 
+
 def parse_juju_origin(data):
     next = _line_generator(data).next
     try:

=== modified file 'juju/providers/common/tests/data/cloud_init_bootstrap'
--- juju/providers/common/tests/data/cloud_init_bootstrap	2012-01-09 14:17:21 +0000
+++ juju/providers/common/tests/data/cloud_init_bootstrap	2012-02-02 16:42:42 +0000
@@ -4,12 +4,58 @@
 machine-data: {juju-provider-type: dummy, juju-zookeeper-hosts: 'localhost:2181',
   machine-id: passport}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper, default-jre-headless, zookeeper, zookeeperd]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'juju-admin initialize --instance-id=token --admin-identity=admin:19vlzY4Vc3q4Ew5OsCwKYqrq1HI= --provider-type=dummy',
-  'JUJU_MACHINE_ID=passport JUJU_ZOOKEEPER=localhost:2181 python -m juju.agents.machine
-    -n --logfile=/var/log/juju/machine-agent.log --pidfile=/var/run/juju/machine-agent.pid',
-  'JUJU_ZOOKEEPER=localhost:2181 python -m juju.agents.provision -n --logfile=/var/log/juju/provision-agent.log
-    --pidfile=/var/run/juju/provision-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper,
+  default-jre-headless, zookeeper, zookeeperd]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'juju-admin initialize --instance-id=token --admin-identity=admin:19vlzY4Vc3q4Ew5OsCwKYqrq1HI=
+    --provider-type=dummy', 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="passport"
+
+    env JUJU_ZOOKEEPER="localhost:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent, 'cat >> /etc/init/juju-provision-agent.conf
+    <<EOF
+
+    description "Juju provisioning agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_ZOOKEEPER="localhost:2181"
+
+
+    exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession >> /tmp/juju-provision-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-provision-agent]
 ssh_authorized_keys: [chubb]

=== modified file 'juju/providers/common/tests/data/cloud_init_bootstrap_zookeepers'
--- juju/providers/common/tests/data/cloud_init_bootstrap_zookeepers	2012-01-09 14:17:21 +0000
+++ juju/providers/common/tests/data/cloud_init_bootstrap_zookeepers	2012-02-02 16:42:42 +0000
@@ -4,13 +4,58 @@
 machine-data: {juju-provider-type: dummy, juju-zookeeper-hosts: 'cotswold:2181,longleat:2181,localhost:2181',
   machine-id: passport}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper, default-jre-headless, zookeeper, zookeeperd]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'juju-admin initialize --instance-id=token --admin-identity=admin:19vlzY4Vc3q4Ew5OsCwKYqrq1HI= --provider-type=dummy',
-  'JUJU_MACHINE_ID=passport JUJU_ZOOKEEPER=cotswold:2181,longleat:2181,localhost:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid', 'JUJU_ZOOKEEPER=cotswold:2181,longleat:2181,localhost:2181
-    python -m juju.agents.provision -n --logfile=/var/log/juju/provision-agent.log
-    --pidfile=/var/run/juju/provision-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper,
+  default-jre-headless, zookeeper, zookeeperd]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'juju-admin initialize --instance-id=token --admin-identity=admin:19vlzY4Vc3q4Ew5OsCwKYqrq1HI=
+    --provider-type=dummy', 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="passport"
+
+    env JUJU_ZOOKEEPER="cotswold:2181,longleat:2181,localhost:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent, 'cat >> /etc/init/juju-provision-agent.conf
+    <<EOF
+
+    description "Juju provisioning agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_ZOOKEEPER="cotswold:2181,longleat:2181,localhost:2181"
+
+
+    exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession >> /tmp/juju-provision-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-provision-agent]
 ssh_authorized_keys: [chubb]

=== modified file 'juju/providers/common/tests/data/cloud_init_branch'
--- juju/providers/common/tests/data/cloud_init_branch	2012-01-09 14:17:21 +0000
+++ juju/providers/common/tests/data/cloud_init_branch	2012-02-02 16:42:42 +0000
@@ -6,12 +6,34 @@
 machine-data: {juju-provider-type: dummy, juju-zookeeper-hosts: 'cotswold:2181,longleat:2181',
   machine-id: passport}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
 runcmd: [sudo apt-get install -y python-txzookeeper, sudo mkdir -p /usr/lib/juju,
-  'cd /usr/lib/juju && sudo /usr/bin/bzr co lp:blah/juju/blah-blah juju',
-  cd /usr/lib/juju/juju && sudo python setup.py develop, sudo mkdir -p /var/lib/juju,
-  sudo mkdir -p /var/log/juju, 'JUJU_MACHINE_ID=passport JUJU_ZOOKEEPER=cotswold:2181,longleat:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+  'cd /usr/lib/juju && sudo /usr/bin/bzr co lp:blah/juju/blah-blah juju', cd /usr/lib/juju/juju
+    && sudo python setup.py develop, sudo mkdir -p /var/lib/juju, sudo mkdir -p /var/log/juju,
+  'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="passport"
+
+    env JUJU_ZOOKEEPER="cotswold:2181,longleat:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [chubb]

=== modified file 'juju/providers/common/tests/data/cloud_init_branch_trunk'
--- juju/providers/common/tests/data/cloud_init_branch_trunk	2012-01-09 14:17:21 +0000
+++ juju/providers/common/tests/data/cloud_init_branch_trunk	2012-02-02 16:42:42 +0000
@@ -6,12 +6,34 @@
 machine-data: {juju-provider-type: dummy, juju-zookeeper-hosts: 'cotswold:2181,longleat:2181',
   machine-id: passport}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
 runcmd: [sudo apt-get install -y python-txzookeeper, sudo mkdir -p /usr/lib/juju,
-  'cd /usr/lib/juju && sudo /usr/bin/bzr co lp:juju juju',
-  cd /usr/lib/juju/juju && sudo python setup.py develop, sudo mkdir -p /var/lib/juju,
-  sudo mkdir -p /var/log/juju, 'JUJU_MACHINE_ID=passport JUJU_ZOOKEEPER=cotswold:2181,longleat:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+  'cd /usr/lib/juju && sudo /usr/bin/bzr co lp:juju juju', cd /usr/lib/juju/juju &&
+    sudo python setup.py develop, sudo mkdir -p /var/lib/juju, sudo mkdir -p /var/log/juju,
+  'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="passport"
+
+    env JUJU_ZOOKEEPER="cotswold:2181,longleat:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [chubb]

=== modified file 'juju/providers/common/tests/data/cloud_init_distro'
--- juju/providers/common/tests/data/cloud_init_distro	2012-01-09 14:17:21 +0000
+++ juju/providers/common/tests/data/cloud_init_distro	2012-02-02 16:42:42 +0000
@@ -4,10 +4,32 @@
 machine-data: {juju-provider-type: dummy, juju-zookeeper-hosts: 'cotswold:2181,longleat:2181',
   machine-id: passport}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'JUJU_MACHINE_ID=passport JUJU_ZOOKEEPER=cotswold:2181,longleat:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="passport"
+
+    env JUJU_ZOOKEEPER="cotswold:2181,longleat:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [chubb]

=== modified file 'juju/providers/common/tests/data/cloud_init_ppa'
--- juju/providers/common/tests/data/cloud_init_ppa	2012-01-09 14:17:21 +0000
+++ juju/providers/common/tests/data/cloud_init_ppa	2012-02-02 16:42:42 +0000
@@ -2,14 +2,36 @@
 apt-update: true
 apt-upgrade: true
 apt_sources:
-- {'source': 'ppa:juju/pkgs'}
+- {source: 'ppa:juju/pkgs'}
 machine-data: {juju-provider-type: dummy, juju-zookeeper-hosts: 'cotswold:2181,longleat:2181',
   machine-id: passport}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'JUJU_MACHINE_ID=passport JUJU_ZOOKEEPER=cotswold:2181,longleat:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="passport"
+
+    env JUJU_ZOOKEEPER="cotswold:2181,longleat:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [chubb]

=== modified file 'juju/providers/common/tests/test_cloudinit.py'
--- juju/providers/common/tests/test_cloudinit.py	2011-10-04 21:22:48 +0000
+++ juju/providers/common/tests/test_cloudinit.py	2012-02-02 16:42:42 +0000
@@ -44,8 +44,9 @@
     def assert_render(self, cloud_init, name):
         with open(os.path.join(DATA_DIR, name)) as f:
             expected = yaml.load(f.read())
-        obtained = yaml.load(cloud_init.render())
-        self.assertEquals(obtained, expected)
+        rendered = cloud_init.render()
+        self.assertTrue(rendered.startswith("#cloud-config"))
+        self.assertEquals(yaml.load(rendered), expected)
 
     def test_render_validate_normal(self):
         cloud_init = CloudInit()

=== modified file 'juju/providers/ec2/tests/data/bootstrap_cloud_init'
--- juju/providers/ec2/tests/data/bootstrap_cloud_init	2012-01-09 14:17:21 +0000
+++ juju/providers/ec2/tests/data/bootstrap_cloud_init	2012-02-02 16:42:42 +0000
@@ -1,16 +1,61 @@
 #cloud-config
 apt-update: true
 apt-upgrade: true
-machine-data: {juju-provider-type: ec2, juju-zookeeper-hosts: 'localhost:2181',
-  machine-id: '0'}
+machine-data: {juju-provider-type: ec2, juju-zookeeper-hosts: 'localhost:2181', machine-id: '0'}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper, default-jre-headless, zookeeper, zookeeperd]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'juju-admin initialize --instance-id=$(curl http://169.254.169.254/1.0/meta-data/instance-id)
-    --admin-identity=admin:JbJ6sDGV37EHzbG9FPvttk64cmg= --provider-type=ec2', 'JUJU_MACHINE_ID=0 JUJU_ZOOKEEPER=localhost:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid', 'JUJU_ZOOKEEPER=localhost:2181
-    python -m juju.agents.provision -n --logfile=/var/log/juju/provision-agent.log
-    --pidfile=/var/run/juju/provision-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper,
+  default-jre-headless, zookeeper, zookeeperd]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'juju-admin initialize --instance-id=$(curl http://169.254.169.254/1.0/meta-data/instance-id)
+    --admin-identity=admin:JbJ6sDGV37EHzbG9FPvttk64cmg= --provider-type=ec2', 'cat
+    >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="0"
+
+    env JUJU_ZOOKEEPER="localhost:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent, 'cat >> /etc/init/juju-provision-agent.conf
+    <<EOF
+
+    description "Juju provisioning agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_ZOOKEEPER="localhost:2181"
+
+
+    exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession >> /tmp/juju-provision-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-provision-agent]
 ssh_authorized_keys: [zebra]

=== modified file 'juju/providers/ec2/tests/data/launch_cloud_init'
--- juju/providers/ec2/tests/data/launch_cloud_init	2012-01-09 14:17:21 +0000
+++ juju/providers/ec2/tests/data/launch_cloud_init	2012-02-02 16:42:42 +0000
@@ -4,10 +4,32 @@
 machine-data: {juju-provider-type: ec2, juju-zookeeper-hosts: 'es.example.internal:2181',
   machine-id: '1'}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'JUJU_MACHINE_ID=1 JUJU_ZOOKEEPER=es.example.internal:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="1"
+
+    env JUJU_ZOOKEEPER="es.example.internal:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [zebra]

=== modified file 'juju/providers/ec2/tests/data/launch_cloud_init_branch'
--- juju/providers/ec2/tests/data/launch_cloud_init_branch	2012-01-09 14:17:21 +0000
+++ juju/providers/ec2/tests/data/launch_cloud_init_branch	2012-02-02 16:42:42 +0000
@@ -6,15 +6,34 @@
 machine-data: {juju-provider-type: ec2, juju-zookeeper-hosts: 'es.example.internal:2181',
   machine-id: '1'}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
-runcmd: [sudo apt-get install -y python-txzookeeper,
-         sudo mkdir -p /usr/lib/juju,
-         'cd /usr/lib/juju && sudo /usr/bin/bzr co lp:~wizard/juju-juicebar juju',
-         cd /usr/lib/juju/juju && sudo python setup.py develop,
-         sudo mkdir -p /var/lib/juju,
-         sudo mkdir -p /var/log/juju,
-         'JUJU_MACHINE_ID=1 JUJU_ZOOKEEPER=es.example.internal:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
+runcmd: [sudo apt-get install -y python-txzookeeper, sudo mkdir -p /usr/lib/juju,
+  'cd /usr/lib/juju && sudo /usr/bin/bzr co lp:~wizard/juju-juicebar juju', cd /usr/lib/juju/juju
+    && sudo python setup.py develop, sudo mkdir -p /var/lib/juju, sudo mkdir -p /var/log/juju,
+  'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="1"
+
+    env JUJU_ZOOKEEPER="es.example.internal:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [zebra]

=== modified file 'juju/providers/ec2/tests/data/launch_cloud_init_ppa'
--- juju/providers/ec2/tests/data/launch_cloud_init_ppa	2012-01-09 14:17:21 +0000
+++ juju/providers/ec2/tests/data/launch_cloud_init_ppa	2012-02-02 16:42:42 +0000
@@ -6,10 +6,32 @@
 machine-data: {juju-provider-type: ec2, juju-zookeeper-hosts: 'es.example.internal:2181',
   machine-id: '1'}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper]
-runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir
-    -p /var/log/juju, 'JUJU_MACHINE_ID=1 JUJU_ZOOKEEPER=es.example.internal:2181
-    python -m juju.agents.machine -n --logfile=/var/log/juju/machine-agent.log
-    --pidfile=/var/run/juju/machine-agent.pid']
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper]
+runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
+    /var/log/juju, 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="1"
+
+    env JUJU_ZOOKEEPER="es.example.internal:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent]
 ssh_authorized_keys: [zebra]

=== modified file 'juju/providers/local/__init__.py'
--- juju/providers/local/__init__.py	2011-11-16 13:56:03 +0000
+++ juju/providers/local/__init__.py	2012-02-02 16:42:42 +0000
@@ -102,11 +102,11 @@
         # Starting provider storage server
         log.info("Starting storage server...")
         storage_server = StorageServer(
-            pid_file=os.path.join(self._directory, "storage-server.pid"),
+            self._qualified_name,
             storage_dir=os.path.join(self._directory, "files"),
             host=net_attributes["ip"]["address"],
             port=get_open_port(net_attributes["ip"]["address"]),
-            log_file=os.path.join(self._directory, "storage-server.log"))
+            logfile=os.path.join(self._directory, "storage-server.log"))
         yield storage_server.start()
 
         # Save the zookeeper start to provider storage.
@@ -130,17 +130,15 @@
             raise ProviderError(str(e))
 
         # Startup the machine agent
-        pid_file = os.path.join(self._directory, "machine-agent.pid")
         log_file = os.path.join(self._directory, "machine-agent.log")
 
         juju_origin = self.config.get("juju-origin")
-        agent = ManagedMachineAgent(pid_file,
+        agent = ManagedMachineAgent(self._qualified_name,
                                     zookeeper_hosts=zookeeper.address,
                                     machine_id="0",
                                     juju_directory=self._directory,
                                     log_file=log_file,
                                     juju_origin=juju_origin,
-                                    juju_unit_namespace=self._qualified_name,
                                     public_key=public_key)
         log.info(
             "Starting machine agent (origin: %s)... ", agent.juju_origin)
@@ -158,14 +156,12 @@
 
         # Stop the machine agent
         log.debug("Stopping machine agent...")
-        pid_file = os.path.join(self._directory, "machine-agent.pid")
-        agent = ManagedMachineAgent(pid_file)
+        agent = ManagedMachineAgent(self._qualified_name)
         yield agent.stop()
 
         # Stop the storage server
         log.debug("Stopping storage server...")
-        pid_file = os.path.join(self._directory, "storage-server.pid")
-        storage_server = StorageServer(pid_file)
+        storage_server = StorageServer(self._qualified_name)
         yield storage_server.stop()
 
         # Stop zookeeper

=== modified file 'juju/providers/local/agent.py'
--- juju/providers/local/agent.py	2011-10-05 12:14:41 +0000
+++ juju/providers/local/agent.py	2012-02-02 16:42:42 +0000
@@ -1,12 +1,6 @@
-import errno
-import os
-import pipes
-import subprocess
 import sys
 
-from twisted.internet.defer import inlineCallbacks, returnValue
-from twisted.internet.threads import deferToThread
-
+from juju.lib.upstart import UpstartService
 from juju.providers.common.cloudinit import get_default_origin, BRANCH
 
 
@@ -15,11 +9,10 @@
     agent_module = "juju.agents.machine"
 
     def __init__(
-        self, pid_file, zookeeper_hosts=None, machine_id="0",
-        log_file=None, juju_directory="/var/lib/juju",
-        juju_unit_namespace="", public_key=None, juju_origin="ppa"):
+            self, juju_unit_namespace, zookeeper_hosts=None,
+            machine_id="0", log_file=None, juju_directory="/var/lib/juju",
+            public_key=None, juju_origin="ppa"):
         """
-        :param pid_file: Path to file used to store process id.
         :param machine_id: machine id for the local machine.
         :param zookeeper_hosts: Zookeeper hosts to connect.
         :param log_file: A file to use for the agent logs.
@@ -31,95 +24,46 @@
         :param public_key: An SSH public key (string) that will be
                used in the container for access.
         """
-        self._pid_file = pid_file
-        self._machine_id = machine_id
-        self._zookeeper_hosts = zookeeper_hosts
-        self._juju_directory = juju_directory
-        self._juju_unit_namespace = juju_unit_namespace
-        self._log_file = log_file
-        self._public_key = public_key
         self._juju_origin = juju_origin
-
         if self._juju_origin is None:
             origin, source = get_default_origin()
             if origin == BRANCH:
                 origin = source
             self._juju_origin = origin
 
+        env = {"JUJU_MACHINE_ID": machine_id,
+               "JUJU_ZOOKEEPER": zookeeper_hosts,
+               "JUJU_HOME": juju_directory,
+               "JUJU_ORIGIN": self._juju_origin,
+               "JUJU_UNIT_NS": juju_unit_namespace,
+               "PYTHONPATH": ":".join(sys.path)}
+        if public_key:
+            env["JUJU_PUBLIC_KEY"] = public_key
+
+        self._service = UpstartService(
+            "juju-%s-machine-agent" % juju_unit_namespace, use_sudo=True)
+        self._service.set_description(
+            "Juju machine agent for %s" % juju_unit_namespace)
+        self._service.set_environ(env)
+        self._service_args = [
+            "/usr/bin/python", "-m", self.agent_module,
+            "--nodaemon", "--logfile", log_file,
+            "--session-file",
+            "/var/run/juju/%s-machine-agent.zksession" % juju_unit_namespace]
+
     @property
     def juju_origin(self):
         return self._juju_origin
 
-    @inlineCallbacks
     def start(self):
-        """Start the machine agent.
-        """
-        assert self._zookeeper_hosts and self._log_file
-
-        if (yield self.is_running()):
-            return
-
-        # sudo even with -E will strip pythonpath, so pass it directly
-        # to the command.
-        args = ["sudo",
-                "JUJU_ZOOKEEPER=%s" % self._zookeeper_hosts,
-                "JUJU_ORIGIN=%s" % self._juju_origin,
-                "JUJU_MACHINE_ID=%s" % self._machine_id,
-                "JUJU_HOME=%s" % self._juju_directory,
-                "JUJU_UNIT_NS=%s" % self._juju_unit_namespace,
-                "PYTHONPATH=%s" % ":".join(sys.path),
-                sys.executable, "-m", self.agent_module,
-                "-n", "--pidfile", self._pid_file,
-                "--logfile", self._log_file]
-
-        if self._public_key:
-            args.insert(
-                1, "JUJU_PUBLIC_KEY=%s" % pipes.quote(self._public_key))
-
-        yield deferToThread(subprocess.check_call, args)
-
-    @inlineCallbacks
+        """Start the machine agent."""
+        self._service.set_command(" ".join(self._service_args))
+        return self._service.start()
+
     def stop(self):
-        """Stop the machine agent.
-        """
-        pid = yield self._get_pid()
-        if pid is None:
-            return
-
-        # Verify the cmdline before attempting to kill.
-        try:
-            with open("/proc/%s/cmdline" % pid) as cmd_file:
-                cmdline = cmd_file.read()
-                if self.agent_module not in cmdline:
-                    raise RuntimeError("Mismatch cmdline")
-        except IOError, e:
-            # Process already died.
-            if e.errno == errno.ENOENT:
-                return
-
-        yield deferToThread(
-            subprocess.check_call, ["sudo", "kill", str(pid)])
-
-    @inlineCallbacks
-    def _get_pid(self):
-        """Return the agent process id or None.
-        """
-        # Default root pidfile mask is restrictive
-        try:
-            pid = yield deferToThread(
-                subprocess.check_output,
-                ["sudo", "cat", self._pid_file],
-                stderr=subprocess.STDOUT)
-        except subprocess.CalledProcessError:
-            return
-        if not pid:
-            return
-        returnValue(int(pid.strip()))
-
-    @inlineCallbacks
+        """Stop the machine agent."""
+        return self._service.destroy()
+
     def is_running(self):
         """Boolean value, true if the machine agent is running."""
-        pid = yield self._get_pid()
-        if pid is None:
-            returnValue(False)
-        returnValue(os.path.isdir("/proc/%s" % pid))
+        return self._service.is_running()

=== modified file 'juju/providers/local/files.py'
--- juju/providers/local/files.py	2011-10-07 18:19:58 +0000
+++ juju/providers/local/files.py	2012-02-02 16:42:42 +0000
@@ -1,13 +1,14 @@
-import errno
+from getpass import getuser
 import os
-import signal
 from StringIO import StringIO
-import subprocess
 import yaml
 
 from twisted.internet.defer import inlineCallbacks, returnValue
+from twisted.internet.error import ConnectionRefusedError
+from twisted.web.client import getPage
 
 from juju.errors import ProviderError, FileNotFound
+from juju.lib.upstart import UpstartService
 from juju.providers.common.files import FileStorage
 
 
@@ -16,22 +17,46 @@
 
 class StorageServer(object):
 
-    def __init__(
-        self, pid_file, storage_dir=None, host=None, port=None, log_file=None):
+    def __init__(self, juju_unit_namespace, storage_dir=None,
+                 host=None, port=None, logfile=None):
         """Management facade for a web server on top of the provider storage.
 
-        :param pid_file: Path to the web server pid file.
+        :param juju_unit_namespace: For disambiguation.
         :param host: Host interface to bind to.
         :param port: Port to bind to.
-        :param log_file: Path to store log output.
+        :param logfile: Path to store log output.
         """
         if storage_dir:
             storage_dir = os.path.abspath(storage_dir)
         self._storage_dir = storage_dir
         self._host = host
         self._port = port
-        self._pid_file = pid_file
-        self._log_file = log_file
+        self._logfile = logfile
+
+        self._service = UpstartService(
+            "juju-%s-file-storage" % juju_unit_namespace, use_sudo=True)
+        self._service.set_description(
+            "Juju file storage for %s" % juju_unit_namespace)
+        self._service_args = [
+            "twistd",
+            "--nodaemon",
+            "--uid", str(os.getuid()),
+            "--gid", str(os.getgid()),
+            "--logfile", logfile,
+            "--pidfile=",
+            "-d", self._storage_dir,
+            "web",
+            "--port", "tcp:%s:interface=%s" % (self._port, self._host),
+            "--path", self._storage_dir]
+
+    @inlineCallbacks
+    def is_serving(self):
+        try:
+            storage = LocalStorage(self._storage_dir)
+            yield getPage((yield storage.get_url(SERVER_URL_KEY)))
+            returnValue(True)
+        except ConnectionRefusedError:
+            returnValue(False)
 
     @inlineCallbacks
     def start(self):
@@ -39,44 +64,33 @@
 
         Also stores the storage server url directly into provider storage.
         """
-        assert (self._storage_dir and self._host
-                and self._port and self._log_file), "Missing start params."
+        assert self._storage_dir, "no storage_dir set"
+        assert self._host, "no host set"
+        assert self._port, "no port set"
+        assert None not in self._service_args, "unset params"
         assert os.path.exists(self._storage_dir), "Invalid storage directory"
+        try:
+            with open(self._logfile, "a"):
+                pass
+        except IOError:
+            raise AssertionError("logfile not writable by this user")
+
 
         storage = LocalStorage(self._storage_dir)
         yield storage.put(
             SERVER_URL_KEY,
             StringIO(yaml.safe_dump(
-                {"storage-url": "http://%s:%s/" % (
-                    self._host, self._port)})))
-
-        subprocess.check_output(
-            ["twistd",
-            "--pidfile", self._pid_file,
-            "--logfile", self._log_file,
-            "-d", self._storage_dir,
-            "web", "--port",
-            "tcp:%s:interface=%s" % (self._port, self._host),
-             "--path", self._storage_dir])
+                {"storage-url": "http://%s:%s/" % (self._host, self._port)})))
+
+        self._service.set_command(" ".join(self._service_args))
+        yield self._service.start()
+
+    def get_pid(self):
+        return self._service.get_pid()
 
     def stop(self):
-        """Stop the storage server.
-        """
-        try:
-            with open(self._pid_file) as pid_file:
-                pid = int(pid_file.read().strip())
-        except IOError:
-            # No pid, move on
-            return
-
-        try:
-            os.kill(pid, 0)
-        except OSError, e:
-            if e.errno == errno.ESRCH:  # No such process, already dead.
-                return
-            raise
-
-        os.kill(pid, signal.SIGKILL)
+        """Stop the storage server."""
+        return self._service.destroy()
 
 
 class LocalStorage(FileStorage):

=== modified file 'juju/providers/local/tests/test_agent.py'
--- juju/providers/local/tests/test_agent.py	2011-10-05 12:14:41 +0000
+++ juju/providers/local/tests/test_agent.py	2012-02-02 16:42:42 +0000
@@ -1,10 +1,13 @@
 import os
 import tempfile
 import subprocess
+import sys
+
 from twisted.internet.defer import inlineCallbacks, succeed
 
+from juju.lib.lxc.tests.test_lxc import uses_sudo
 from juju.lib.testing import TestCase
-from juju.lib.lxc.tests.test_lxc import run_lxc_tests
+from juju.lib.upstart import UpstartService
 from juju.tests.common import get_test_zookeeper_address
 from juju.providers.local.agent import ManagedMachineAgent
 
@@ -12,61 +15,93 @@
 class ManagedAgentTest(TestCase):
 
     @inlineCallbacks
-    def test_managed_agent_args(self):
-
-        captured_args = []
-
-        def intercept_args(args):
-            captured_args.extend(args)
-            return True
-
-        self.patch(subprocess, "check_call", intercept_args)
+    def test_managed_agent_config(self):
+        subprocess_calls = []
+
+        def intercept_args(args, **kwargs):
+            subprocess_calls.append(args)
+            self.assertEquals(args[0], "sudo")
+            if args[1] == "cp":
+                return real_check_call(args[1:], **kwargs)
+            return 0
+
+        real_check_call = self.patch(subprocess, "check_call", intercept_args)
+        init_dir = self.makeDir()
+        self.patch(UpstartService, "init_dir", init_dir)
+
+        # Mock out the repeated checking for unstable pid, after an initial
+        # stop/waiting to induce the actual start
+        getProcessOutput = self.mocker.replace(
+            "twisted.internet.utils.getProcessOutput")
+        getProcessOutput("/sbin/status", ["juju-ns1-machine-agent"])
+        self.mocker.result(succeed("stop/waiting"))
+        for _ in range(5):
+            getProcessOutput("/sbin/status", ["juju-ns1-machine-agent"])
+            self.mocker.result(succeed("start/running 123"))
+        self.mocker.replay()
 
         juju_directory = self.makeDir()
-        pid_file = self.makeFile()
         log_file = self.makeFile()
-
         agent = ManagedMachineAgent(
-            pid_file, get_test_zookeeper_address(),
+            "ns1",
+            get_test_zookeeper_address(),
             juju_directory=juju_directory,
-            log_file=log_file, juju_unit_namespace="ns1",
+            log_file=log_file,
             juju_origin="lp:juju/trunk")
 
-        mock_agent = self.mocker.patch(agent)
-        mock_agent.is_running()
-        self.mocker.result(succeed(False))
-        self.mocker.replay()
+        try:
+            os.remove("/tmp/juju-ns1-machine-agent.output")
+        except OSError:
+            pass  # just make sure it's not there, so the .start()
+                  # doesn't insert a spurious rm
 
-        self.assertEqual(agent.juju_origin, "lp:juju/trunk")
         yield agent.start()
 
-        # Verify machine agent environment
-        env_vars = dict(
-            [arg.split("=") for arg in captured_args if "=" in arg])
-        env_vars.pop("PYTHONPATH")
-        self.assertEqual(
-            env_vars,
-            dict(JUJU_ZOOKEEPER=get_test_zookeeper_address(),
-                 JUJU_MACHINE_ID="0",
-                 JUJU_HOME=juju_directory,
-                 JUJU_ORIGIN="lp:juju/trunk",
-                 JUJU_UNIT_NS="ns1"))
-
+        conf_dest = os.path.join(
+            init_dir, "juju-ns1-machine-agent.conf")
+        chmod, start = subprocess_calls[1:]
+        self.assertEquals(chmod, ("sudo", "chmod", "644", conf_dest))
+        self.assertEquals(
+            start, ("sudo", "/sbin/start", "juju-ns1-machine-agent"))
+
+        env = []
+        with open(conf_dest) as f:
+            for line in f:
+                if line.startswith("env"):
+                    env.append(line[4:-1].split("=", 1))
+                if line.startswith("exec"):
+                    exec_ = line[5:-1]
+
+        expect_exec = (
+            "/usr/bin/python -m juju.agents.machine --nodaemon --logfile %s "
+            "--session-file /var/run/juju/ns1-machine-agent.zksession "
+            ">> /tmp/juju-ns1-machine-agent.output 2>&1"
+            % log_file)
+        self.assertEquals(exec_, expect_exec)
+
+        env = dict((k, v.strip('"')) for (k, v) in env)
+        self.assertEquals(env, {
+            "JUJU_ZOOKEEPER": get_test_zookeeper_address(),
+            "JUJU_MACHINE_ID": "0",
+            "JUJU_HOME": juju_directory,
+            "JUJU_ORIGIN": "lp:juju/trunk",
+            "JUJU_UNIT_NS": "ns1",
+            "PYTHONPATH": ":".join(sys.path)})
+
+    @uses_sudo
     @inlineCallbacks
     def test_managed_agent_root(self):
         juju_directory = self.makeDir()
-        pid_file = tempfile.mktemp()
         log_file = tempfile.mktemp()
 
         # The pid file and log file get written as root
         def cleanup_root_file(cleanup_file):
             subprocess.check_call(
                 ["sudo", "rm", "-f", cleanup_file], stderr=subprocess.STDOUT)
-        self.addCleanup(cleanup_root_file, pid_file)
         self.addCleanup(cleanup_root_file, log_file)
 
         agent = ManagedMachineAgent(
-            pid_file, machine_id="0", log_file=log_file,
+            "test-ns", machine_id="0", log_file=log_file,
             zookeeper_hosts=get_test_zookeeper_address(),
             juju_directory=juju_directory)
 
@@ -85,13 +120,3 @@
 
         # running stop again is fine, detects the process is stopped.
         yield agent.stop()
-
-        self.assertFalse(os.path.exists(pid_file))
-
-        # Stop raises runtime error if the process doesn't match up.
-        with open(pid_file, "w") as pid_handle:
-            pid_handle.write("1")
-        self.assertFailure(agent.stop(), RuntimeError)
-
-    # Reuse the lxc flag for tests needing sudo
-    test_managed_agent_root.skip = run_lxc_tests()

=== modified file 'juju/providers/local/tests/test_files.py'
--- juju/providers/local/tests/test_files.py	2011-10-07 18:19:58 +0000
+++ juju/providers/local/tests/test_files.py	2012-02-02 16:42:42 +0000
@@ -1,14 +1,16 @@
 import os
+import signal
 from StringIO  import StringIO
+import subprocess
 import yaml
 
-from twisted.internet.defer import inlineCallbacks
+from twisted.internet.defer import inlineCallbacks, succeed
 from twisted.web.client import getPage
 
+from juju.errors import ProviderError, ServiceError
+from juju.lib.lxc.tests.test_lxc import uses_sudo
 from juju.lib.testing import TestCase
-
-
-from juju.errors import ProviderError
+from juju.lib.upstart import UpstartService
 from juju.providers.local.files import (
     LocalStorage, StorageServer, SERVER_URL_KEY)
 from juju.state.utils import get_open_port
@@ -16,38 +18,151 @@
 
 class WebFileStorageTest(TestCase):
 
+    @inlineCallbacks
     def setUp(self):
+        yield super(WebFileStorageTest, self).setUp()
         self._storage_path = self.makeDir()
+        self._logfile = self.makeFile()
         self._storage = LocalStorage(self._storage_path)
-        self._log_path = self.makeFile()
-        self._pid_path = self.makeFile()
         self._port = get_open_port()
         self._server = StorageServer(
-            self._pid_path, self._storage_path, "localhost",
-            get_open_port(), self._log_path)
+            "ns1", self._storage_path, "localhost", self._port, self._logfile)
 
     @inlineCallbacks
-    def test_start_stop(self):
-        yield self._storage.put("abc", StringIO("hello world"))
-        yield self._server.start()
-        storage_url = yield self._storage.get_url("abc")
-        contents = yield getPage(storage_url)
-        self.assertEqual("hello world", contents)
-        self._server.stop()
-        # Stopping multiple times is fine.
-        self._server.stop()
+    def wait_for_server(self, server):
+        while not (yield server.is_serving()):
+            yield self.sleep(0.1)
 
     def test_start_missing_args(self):
-        server = StorageServer(self._pid_path)
+        server = StorageServer("ns1", self._storage_path)
         return self.assertFailure(server.start(), AssertionError)
 
     def test_start_invalid_directory(self):
         os.rmdir(self._storage_path)
         return self.assertFailure(self._server.start(), AssertionError)
 
-    def test_stop_missing_pid(self):
-        server = StorageServer(self._pid_path)
-        server.stop()
+    @inlineCallbacks
+    def test_upstart(self):
+        subprocess_calls = []
+
+        def intercept_args(args, **kwargs):
+            subprocess_calls.append(args)
+            self.assertEquals(args[0], "sudo")
+            if args[1] == "cp":
+                return real_check_call(args[1:], **kwargs)
+            return 0
+
+        real_check_call = self.patch(subprocess, "check_call", intercept_args)
+        init_dir = self.makeDir()
+        self.patch(UpstartService, "init_dir", init_dir)
+
+        # Mock out the repeated checking for unstable pid, after an initial
+        # stop/waiting to induce the actual start
+        getProcessOutput = self.mocker.replace(
+            "twisted.internet.utils.getProcessOutput")
+        getProcessOutput("/sbin/status", ["juju-ns1-file-storage"])
+        self.mocker.result(succeed("stop/waiting"))
+        for _ in range(5):
+            getProcessOutput("/sbin/status", ["juju-ns1-file-storage"])
+            self.mocker.result(succeed("start/running 123"))
+        self.mocker.replay()
+
+        try:
+            os.remove("/tmp/juju-ns1-file-storage.output")
+        except OSError:
+            pass  # just make sure it's not there, so the .start()
+                  # doesn't insert a spurious rm
+
+        yield self._server.start()
+        chmod = subprocess_calls[1]
+        conf_dest = os.path.join(init_dir, "juju-ns1-file-storage.conf")
+        self.assertEquals(chmod, ("sudo", "chmod", "644", conf_dest))
+        start = subprocess_calls[-1]
+        self.assertEquals(
+            start, ("sudo", "/sbin/start", "juju-ns1-file-storage"))
+
+        with open(conf_dest) as f:
+            for line in f:
+                if line.startswith("env"):
+                    self.fail("didn't expect any special environment")
+                if line.startswith("exec"):
+                    exec_ = line[5:].strip()
+
+        expect_exec = (
+            "twistd --nodaemon --uid %s --gid %s --logfile %s --pidfile= -d "
+            "%s web --port tcp:%s:interface=localhost --path %s >> "
+            "/tmp/juju-ns1-file-storage.output 2>&1"
+            % (os.getuid(), os.getgid(), self._logfile, self._storage_path,
+            self._port, self._storage_path))
+        self.assertEquals(exec_, expect_exec)
+
+    @uses_sudo
+    @inlineCallbacks
+    def test_start_stop(self):
+        yield self._storage.put("abc", StringIO("hello world"))
+        yield self._server.start()
+        # Starting multiple times is fine.
+        yield self._server.start()
+        storage_url = yield self._storage.get_url("abc")
+
+        # It might not have started actually accepting connections yet...
+        yield self.wait_for_server(self._server)
+        self.assertEqual((yield getPage(storage_url)), "hello world")
+
+        # Check that it can be killed by the current user (ie, is not running
+        # as root) and still comes back up
+        old_pid = yield self._server.get_pid()
+        os.kill(old_pid, signal.SIGKILL)
+        new_pid = yield self._server.get_pid()
+        self.assertNotEquals(old_pid, new_pid)
+
+        # Give it a moment to actually start serving again
+        yield self.wait_for_server(self._server)
+        self.assertEqual((yield getPage(storage_url)), "hello world")
+
+        yield self._server.stop()
+        # Stopping multiple times is fine too.
+        yield self._server.stop()
+
+    @uses_sudo
+    @inlineCallbacks
+    def test_namespacing(self):
+        alt_storage_path = self.makeDir()
+        alt_storage = LocalStorage(alt_storage_path)
+        yield alt_storage.put("some-path", StringIO("alternative"))
+        yield self._storage.put("some-path", StringIO("original"))
+
+        alt_server = StorageServer(
+            "ns2", alt_storage_path, "localhost", get_open_port(),
+            self.makeFile())
+        yield alt_server.start()
+        yield self._server.start()
+        yield self.wait_for_server(alt_server)
+        yield self.wait_for_server(self._server)
+
+        alt_contents = yield getPage(
+            (yield alt_storage.get_url("some-path")))
+        self.assertEquals(alt_contents, "alternative")
+        orig_contents = yield getPage(
+            (yield self._storage.get_url("some-path")))
+        self.assertEquals(orig_contents, "original")
+
+        yield alt_server.stop()
+        yield self._server.stop()
+
+    @uses_sudo
+    @inlineCallbacks
+    def test_capture_errors(self):
+        self._port = get_open_port()
+        self._server = StorageServer(
+            "borken", self._storage_path, "lol borken", self._port,
+            self._logfile)
+        d = self._server.start()
+        e = yield self.assertFailure(d, ServiceError)
+        self.assertTrue(str(e).startswith(
+            "Failed to start job juju-borken-file-storage; got output:\n"))
+        self.assertIn("Wrong number of arguments", str(e))
+        yield self._server.stop()
 
 
 class FileStorageTest(TestCase):

=== modified file 'juju/providers/orchestra/tests/data/bootstrap_user_data'
--- juju/providers/orchestra/tests/data/bootstrap_user_data	2012-01-09 14:17:21 +0000
+++ juju/providers/orchestra/tests/data/bootstrap_user_data	2012-02-02 16:42:42 +0000
@@ -1,15 +1,60 @@
-#cloud-config
 apt-update: true
 apt-upgrade: true
 machine-data: {juju-provider-type: orchestra, juju-zookeeper-hosts: 'localhost:2181',
   machine-id: '0'}
 output: {all: '| tee -a /var/log/cloud-init-output.log'}
-packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws,
-  python-zookeeper, default-jre-headless, zookeeper, zookeeperd]
+packages: [bzr, byobu, tmux, python-setuptools, python-twisted, python-txaws, python-zookeeper,
+  default-jre-headless, zookeeper, zookeeperd]
 runcmd: [sudo apt-get -y install juju, sudo mkdir -p /var/lib/juju, sudo mkdir -p
-    /var/log/juju, 'juju-admin initialize --instance-id=winston-uid --admin-identity=admin:qRBXC1ubEEUqRL6wcBhgmc9xkaY= --provider-type=orchestra',
-  'JUJU_MACHINE_ID=0 JUJU_ZOOKEEPER=localhost:2181 python -m juju.agents.machine -n
-    --logfile=/var/log/juju/machine-agent.log --pidfile=/var/run/juju/machine-agent.pid',
-  'JUJU_ZOOKEEPER=localhost:2181 python -m juju.agents.provision -n --logfile=/var/log/juju/provision-agent.log
-    --pidfile=/var/run/juju/provision-agent.pid']
+    /var/log/juju, 'juju-admin initialize --instance-id=winston-uid --admin-identity=admin:qRBXC1ubEEUqRL6wcBhgmc9xkaY=
+    --provider-type=orchestra', 'cat >> /etc/init/juju-machine-agent.conf <<EOF
+
+    description "Juju machine agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_MACHINE_ID="0"
+
+    env JUJU_ZOOKEEPER="localhost:2181"
+
+
+    exec python -m juju.agents.machine --nodaemon --logfile /var/log/juju/machine-agent.log
+    --session-file /var/run/juju/machine-agent.zksession >> /tmp/juju-machine-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-machine-agent, 'cat >> /etc/init/juju-provision-agent.conf
+    <<EOF
+
+    description "Juju provisioning agent"
+
+    author "Juju Team <[e-mail]>"
+
+
+    start on runlevel [2345]
+
+    stop on runlevel [!2345]
+
+    respawn
+
+
+    env JUJU_ZOOKEEPER="localhost:2181"
+
+
+    exec python -m juju.agents.provision --nodaemon --logfile /var/log/juju/provision-agent.log
+    --session-file /var/run/juju/provision-agent.zksession >> /tmp/juju-provision-agent.output
+    2>&1
+
+    EOF
+
+    ', /sbin/start juju-provision-agent]
 ssh_authorized_keys: [this-is-a-public-key]

=== modified file 'juju/providers/orchestra/tests/data/launch_user_data'
--- juju/providers/orchestra/tests/data/launch_user_data	2012-01-09 14:17:21 +0000
+++ juju/providers/orchestra/tests/data/launch_user_data	2012-02-02 16:42:42 +0000
@@ -1,12 +1,34 @@
-#cloud-config
 apt-update: true