From fb146320cb963cefe88fa0a2f4d2358aefe29e57 Mon Sep 17 00:00:00 2001
From: Menno Smits <menno.smits@canonical.com>
Date: Mon, 8 Sep 2014 13:58:02 +1200
Subject: [PATCH] Add new wrench points for upgrade-steps worker

---
 cmd/jujud/upgrade.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/cmd/jujud/upgrade.go b/cmd/jujud/upgrade.go
index 35d6b5e99ea..e9e5dd44400 100644
--- a/cmd/jujud/upgrade.go
+++ b/cmd/jujud/upgrade.go
@@ -17,6 +17,7 @@ import (
 	"github.com/juju/juju/upgrades"
 	"github.com/juju/juju/version"
 	"github.com/juju/juju/worker"
+	"github.com/juju/juju/wrench"
 )
 
 type upgradingMachineAgent interface {
@@ -67,6 +68,11 @@ type upgradeWorkerContext struct {
 // InitialiseUsingAgent sets up a upgradeWorkerContext from a machine agent instance.
 // It may update the agent's configuration.
 func (c *upgradeWorkerContext) InitializeUsingAgent(a upgradingMachineAgent) error {
+	if wrench.IsActive("machine-agent", "always-try-upgrade") {
+		// Always enter upgrade mode. This allows test of upgrades
+		// even when there's actually no upgrade steps to run.
+		return nil
+	}
 	return a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
 		if !upgrades.AreUpgradesDefined(agentConfig.UpgradedToVersion()) {
 			logger.Infof("no upgrade steps required or upgrade steps for %v "+
@@ -115,6 +121,10 @@ func isAPILostDuringUpgrade(err error) bool {
 }
 
 func (c *upgradeWorkerContext) run(stop <-chan struct{}) error {
+	if wrench.IsActive("machine-agent", "fail-upgrade-start") {
+		return nil // Make the worker stop
+	}
+
 	select {
 	case <-c.UpgradeComplete:
 		// Our work is already done (we're probably being restarted
@@ -323,6 +333,14 @@ func (c *upgradeWorkerContext) waitForOtherStateServers(info *state.UpgradeInfo,
 }
 
 func getUpgradeStartTimeout(isMaster bool) time.Duration {
+	if wrench.IsActive("machine-agent", "short-upgrade-timeout") {
+		// This duration is fairly arbitrary. During manual testing it
+		// avoids the normal long wait but still provides a small
+		// window to check the environment status and logs before the
+		// timeout is triggered.
+		return time.Minute
+	}
+
 	if isMaster {
 		return upgradeStartTimeoutMaster
 	}