Merge workloads.go into payloads.go.

nammn · Nov 4, 2015 · dd185ee · dd185ee
1 parent dbc6eb9
commit dd185ee
Show file tree

Hide file tree

Showing 5 changed files with 148 additions and 183 deletions.
diff --git a/component/all/payload.go b/component/all/payload.go
@@ -194,7 +194,6 @@ func (payloads) registerState() {
 	newUnitPayloads := func(persist state.Persistence, unit, machine string) (state.UnitPayloads, error) {
 		return payloadstate.NewUnitPayloads(persist, unit, machine), nil
 	}
-	state.SetWorkloadsComponent(newUnitPayloads)
 
 	newEnvPayloads := func(persist state.PayloadsEnvPersistence) (state.EnvPayloads, error) {
 		envPersist := persistence.NewEnvPersistence(persist)
@@ -203,5 +202,6 @@ func (payloads) registerState() {
 		}
 		return envPayloads, nil
 	}
-	state.SetPayloadsComponent(newEnvPayloads)
+
+	state.SetPayloadsComponent(newEnvPayloads, newUnitPayloads)
 }
diff --git a/state/payloads.go b/state/payloads.go
@@ -18,6 +18,31 @@ type EnvPayloads interface {
 	ListAll() ([]workload.FullPayloadInfo, error)
 }
 
+// UnitPayloads exposes high-level interaction with workloads for a unit.
+type UnitPayloads interface {
+	// Track tracks a workload in state. If a workload is
+	// already being tracked for the same (unit, workload name, plugin ID)
+	// then the request will fail. The unit must also be "alive".
+	Track(payload workload.Payload) error
+	// SetStatus sets the status of a workload. Only some fields
+	// must be set on the provided info: Name, Status, Details.ID, and
+	// Details.Status. If the workload is not in state then the request
+	// will fail.
+	SetStatus(docID, status string) error
+	// List builds the list of workloads registered for
+	// the given unit and IDs. If no IDs are provided then all
+	// registered workloads for the unit are returned. In the case that
+	// IDs are provided, any that are not in state are ignored and only
+	// the found ones are returned. It is up to the caller to
+	// extrapolate the list of missing IDs.
+	List(ids ...string) ([]workload.Result, error)
+	// LookUpReturns the Juju ID for the corresponding workload.
+	LookUp(name, rawID string) (string, error)
+	// Untrack removes the identified workload from state. If the
+	// given ID is not in state then the request will fail.
+	Untrack(id string) error
+}
+
 // TODO(ericsnow) Use a more generic component registration mechanism?
 
 // PayloadsEnvPersistence provides all the information needed to produce
@@ -37,17 +62,19 @@ type PayloadsEnvPersistence interface {
 }
 
 type newEnvPayloadsFunc func(PayloadsEnvPersistence) (EnvPayloads, error)
+type newUnitPayloadsFunc func(persist Persistence, unit, machine string) (UnitPayloads, error)
 
+// TODO(ericsnow) Merge the 2 vars
 var (
-	newEnvPayloads newEnvPayloadsFunc
+	newEnvPayloads  newEnvPayloadsFunc
+	newUnitPayloads newUnitPayloadsFunc
 )
 
-// TODO(ericsnow) Merge the 2 Set*Component funcs.
-
 // SetPayloadComponent registers the functions that provide the state
 // functionality related to payloads.
-func SetPayloadsComponent(epFunc newEnvPayloadsFunc) {
+func SetPayloadsComponent(epFunc newEnvPayloadsFunc, upFunc newUnitPayloadsFunc) {
 	newEnvPayloads = epFunc
+	newUnitPayloads = upFunc
 }
 
 // EnvPayloads exposes interaction with payloads in state.
@@ -68,6 +95,28 @@ func (st *State) EnvPayloads() (EnvPayloads, error) {
 	return envPayloads, nil
 }
 
+// UnitPayloads exposes interaction with workloads in state
+// for a the given unit.
+func (st *State) UnitPayloads(unit *Unit) (UnitPayloads, error) {
+	if newUnitPayloads == nil {
+		return nil, errors.Errorf("payloads not supported")
+	}
+
+	machineID, err := unit.AssignedMachineId()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	unitID := unit.UnitTag().Id()
+
+	persist := st.newPersistence()
+	unitWorkloads, err := newUnitPayloads(persist, unitID, machineID)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	return unitWorkloads, nil
+}
+
 type payloadsEnvPersistence struct {
 	Persistence
 	st *State

diff --git a/state/payloads_test.go b/state/payloads_test.go
@@ -4,6 +4,7 @@
 package state_test
 
 import (
+	"github.com/juju/errors"
 	jc "github.com/juju/testing/checkers"
 	gc "gopkg.in/check.v1"
 	"gopkg.in/juju/charm.v5"
@@ -19,7 +20,10 @@ func init() {
 	}
 }
 
-var _ = gc.Suite(&envPayloadsSuite{})
+var (
+	_ = gc.Suite(&envPayloadsSuite{})
+	_ = gc.Suite(&unitPayloadsSuite{})
+)
 
 type envPayloadsSuite struct {
 	ConnSuite
@@ -84,6 +88,94 @@ func (s *envPayloadsSuite) TestFunctional(c *gc.C) {
 	c.Check(payloads, gc.HasLen, 0)
 }
 
+type unitPayloadsSuite struct {
+	ConnSuite
+}
+
+func (s *unitPayloadsSuite) TestFunctional(c *gc.C) {
+	unit := addUnit(c, s.ConnSuite, unitArgs{
+		charm:    "dummy",
+		service:  "a-service",
+		metadata: payloadsMetaYAML,
+	})
+
+	st, err := s.State.UnitPayloads(unit)
+	c.Assert(err, jc.ErrorIsNil)
+
+	results, err := st.List()
+	c.Assert(err, jc.ErrorIsNil)
+	c.Check(results, gc.HasLen, 0)
+
+	pl := workload.Payload{
+		PayloadClass: charm.PayloadClass{
+			Name: "payloadA",
+			Type: "docker",
+		},
+		ID:     "xyz",
+		Status: workload.StateRunning,
+		Unit:   "a-service/0",
+	}
+	err = st.Track(pl)
+	c.Assert(err, jc.ErrorIsNil)
+
+	results, err = st.List()
+	c.Assert(err, jc.ErrorIsNil)
+	// TODO(ericsnow) Once Track returns the new ID we can drop
+	// the following two lines.
+	c.Assert(results, gc.HasLen, 1)
+	id := results[0].ID
+	c.Check(results, jc.DeepEquals, []workload.Result{{
+		ID: id,
+		Payload: &workload.FullPayloadInfo{
+			Payload: pl,
+			// TODO(ericsnow) Require machine "1"?
+			Machine: "",
+		},
+	}})
+
+	lookedUpID, err := st.LookUp("payloadA", "xyz")
+	c.Assert(err, jc.ErrorIsNil)
+	c.Check(lookedUpID, gc.Equals, id)
+
+	c.Logf("using ID %q", id)
+	results, err = st.List(id)
+	c.Assert(err, jc.ErrorIsNil)
+	c.Check(results, jc.DeepEquals, []workload.Result{{
+		ID: id,
+		Payload: &workload.FullPayloadInfo{
+			Payload: pl,
+			Machine: "",
+		},
+	}})
+
+	err = st.SetStatus(id, "running")
+	c.Assert(err, jc.ErrorIsNil)
+
+	results, err = st.List(id)
+	c.Assert(err, jc.ErrorIsNil)
+	c.Check(results, jc.DeepEquals, []workload.Result{{
+		ID: id,
+		Payload: &workload.FullPayloadInfo{
+			Payload: pl,
+			Machine: "",
+		},
+	}})
+
+	// Ensure duplicates are not allowed.
+	err = st.Track(pl)
+	c.Check(err, jc.Satisfies, errors.IsAlreadyExists)
+	results, err = st.List()
+	c.Assert(err, jc.ErrorIsNil)
+	c.Check(results, gc.HasLen, 1)
+
+	err = st.Untrack(id)
+	c.Assert(err, jc.ErrorIsNil)
+
+	results, err = st.List()
+	c.Assert(err, jc.ErrorIsNil)
+	c.Check(results, gc.HasLen, 0)
+}
+
 const payloadsMetaYAML = `
 name: a-charm
 summary: a charm...

diff --git a/state/workloads.go b/state/workloads.go
diff --git a/state/workloads_test.go b/state/workloads_test.go