Payload fix unregister

workload/context/unregister.go

@@ -0,0 +1,76 @@
+// Copyright 2015 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package context
+
+import (
+	"github.com/juju/cmd"
+	"github.com/juju/errors"
+)
+
+const UntrackCmdName = "payload-unregister"
+
+// NewUntrackCmd returns a new UntrackCmd that wraps the given context.
+func NewUntrackCmd(ctx HookContext) (*UntrackCmd, error) {
+	compCtx, err := ContextComponent(ctx)
+	if err != nil {
+		// The component wasn't tracked properly.
+		return nil, errors.Trace(err)
+	}
+	return &UntrackCmd{api: compCtx}, nil
+}
+
+// Info implements cmd.Command.
+func (c UntrackCmd) Info() *cmd.Info {
+	return &cmd.Info{
+		Name:    "payload-unregister",
+		Args:    "<class> <id>",
+		Purpose: "stop tracking a payload",
+		Doc: `
+"payload-unregister" is used while a hook is running to let Juju know
+that a payload has been manually stopped. The <class> and <id> provided
+must match a payload that has been previously registered with juju using
+payload-register.
+`,
+	}
+}
+
+// UntrackCmd implements the untrack command.
+type UntrackCmd struct {
+	*cmd.CommandBase
+
+	api   Component
+	class string
+	id    string
+}
+
+// Init implements cmd.Command.
+func (c *UntrackCmd) Init(args []string) error {
+	if len(args) < 2 {
+		return errors.Errorf("missing required arguments")
+	}
+	c.class = args[0]
+	c.id = args[1]
+	return cmd.CheckEmpty(args[2:])
+}
+
+// Run runs the untrack command.
+func (c *UntrackCmd) Run(ctx *cmd.Context) error {
+	//TODO(wwitzel3) make Untrack accept class and id and
+	// compose the ID in the API layer using BuildID
+
+	ID := c.class + "/" + c.id
+	logger.Tracef("Running untrack command for %q", ID)
+
+	if err := c.api.Untrack(ID); err != nil {
+		return errors.Trace(err)
+	}
+
+	// We flush to state immedeiately so that status reflects the
+	// workload correctly.
+	if err := c.api.Flush(); err != nil {
+		return errors.Trace(err)
+	}
+
+	return nil
+}

workload/context/untrack_test.go → workload/context/unregister_test.go

@@ -33,19 +33,20 @@ func (s *untrackSuite) TestCommandRegistered(c *gc.C) {
 
 func (s *untrackSuite) TestHelp(c *gc.C) {
 	s.checkHelp(c, `
-usage: payload-unregister <name-or-id>
 purpose: stop tracking a payload
 
 "payload-unregister" is used while a hook is running to let Juju know
-that a payload has been manually stopped. The id
-used to start tracking the payload must be provided.
+payload-register.
 `[1:])
 }
 
 func (s *untrackSuite) TestRunOkay(c *gc.C) {
 	s.setMetadata(s.workload)
 	s.compCtx.workloads[s.workload.ID()] = s.workload
-	err := s.cmd.Init([]string{s.workload.Name})
 	c.Assert(err, jc.ErrorIsNil)
 
 	c.Logf("%#v", s.cmd)

workload/context/untrack.go

@@ -1,49 +0,0 @@
-// Copyright 2015 Canonical Ltd.
-// Licensed under the AGPLv3, see LICENCE file for details.
-
-package context
-
-import (
-	"github.com/juju/cmd"
-	"github.com/juju/errors"
-)
-
-const UntrackCmdName = "payload-unregister"
-
-// NewUntrackCmd returns an UntrackCmd that uses the given hook context.
-func NewUntrackCmd(ctx HookContext) (*UntrackCmd, error) {
-	base, err := newCommand(ctx)
-	if err != nil {
-		return nil, errors.Trace(err)
-	}
-	c := &UntrackCmd{
-		baseCommand: base,
-	}
-	c.cmdInfo = cmdInfo{
-		Name:    "payload-unregister",
-		Summary: "stop tracking a payload",
-		Doc: `
-"payload-unregister" is used while a hook is running to let Juju know
-that a payload has been manually stopped. The id
-used to start tracking the payload must be provided.
-`,
-	}
-	return c, nil
-}
-
-var _ cmd.Command = (*UntrackCmd)(nil)
-
-// UntrackCmd implements the untrack command.
-type UntrackCmd struct {
-	*baseCommand
-}
-
-// Run runs the untrack command.
-func (c *UntrackCmd) Run(ctx *cmd.Context) error {
-	logger.Tracef("Running untrack command with id %q", c.ID)
-	if err := c.baseCommand.Run(ctx); err != nil {
-		return errors.Trace(err)
-	}
-
-	return c.compCtx.Untrack(c.ID)
-}