use class/id for untrack

workload/context/unregister.go

@@ -10,40 +10,64 @@ import (
 
 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 {
+		// The component wasn't tracked properly.
 		return nil, errors.Trace(err)
 	}
-	c := &UntrackCmd{
-		baseCommand: base,
-	}
-	c.cmdInfo = cmdInfo{
+func (c UntrackCmd) Info() *cmd.Info {
+	return &cmd.Info{
 		Name:    "payload-unregister",
-		Summary: "stop tracking a payload",
+		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 id
-used to start tracking the payload must be provided.
+payload-register.
 `,
 	}
-	return c, nil
 }
 
-var _ cmd.Command = (*UntrackCmd)(nil)
-
 // UntrackCmd implements the untrack command.
 type UntrackCmd struct {
-	*baseCommand
+
+	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 nil
 }
 
 // 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 {
+
+	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 c.compCtx.Untrack(c.ID)
 }

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)