Add a Machine field to UnitPayloads.

nammn · Nov 4, 2015 · a6226f9 · a6226f9
1 parent 8dcd685
commit a6226f9
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 7 deletions.
diff --git a/component/all/payload.go b/component/all/payload.go
@@ -191,8 +191,8 @@ func (payloads) registerState() {
 	// TODO(ericsnow) Use a more general registration mechanism.
 	//state.RegisterMultiEnvCollections(persistence.Collections...)
 
-	newUnitPayloads := func(persist state.Persistence, unit string) (state.UnitPayloads, error) {
-		return payloadstate.NewUnitPayloads(persist, unit), nil
+	newUnitPayloads := func(persist state.Persistence, unit, machine string) (state.UnitPayloads, error) {
+		return payloadstate.NewUnitPayloads(persist, unit, machine), nil
 	}
 	state.SetWorkloadsComponent(newUnitPayloads)
 

diff --git a/state/workloads.go b/state/workloads.go
@@ -38,7 +38,7 @@ type UnitPayloads interface {
 
 // TODO(ericsnow) Use a more generic component registration mechanism?
 
-type newUnitPayloadsFunc func(persist Persistence, unit string) (UnitPayloads, error)
+type newUnitPayloadsFunc func(persist Persistence, unit, machine string) (UnitPayloads, error)
 
 var (
 	newUnitPayloads newUnitPayloadsFunc
@@ -57,8 +57,14 @@ func (st *State) UnitPayloads(unit *Unit) (UnitPayloads, error) {
 		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, unit.UnitTag().Id())
+	unitWorkloads, err := newUnitPayloads(persist, unitID, machineID)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

diff --git a/workload/state/unit.go b/workload/state/unit.go
@@ -36,15 +36,19 @@ type UnitPayloads struct {
 	// Unit identifies the unit associated with the workloads.
 	Unit string
 
+	// Machine identifies the unit's machine.
+	Machine string
+
 	newID func() (string, error)
 }
 
 // NewUnitPayloads builds a UnitPayloads for a unit.
-func NewUnitPayloads(st persistence.PersistenceBase, unit string) *UnitPayloads {
+func NewUnitPayloads(st persistence.PersistenceBase, unit, machine string) *UnitPayloads {
 	persist := persistence.NewPersistence(st, unit)
 	return &UnitPayloads{
 		Persist: persist,
 		Unit:    unit,
+		Machine: machine,
 		newID:   newID,
 	}
 }
@@ -145,9 +149,14 @@ func (uw UnitPayloads) List(ids ...string) ([]workload.Result, error) {
 		pl := payloads[i]
 		i += 1
 
+		// TODO(ericsnow) Ensure that pl.Unit == uw.Unit?
+
 		result := workload.Result{
-			ID:      id,
-			Payload: &workload.FullPayloadInfo{Payload: pl},
+			ID: id,
+			Payload: &workload.FullPayloadInfo{
+				Payload: pl,
+				Machine: uw.Machine,
+			},
 		}
 		if id == "" {
 			// TODO(ericsnow) Do this more efficiently.