Merge unit and env persistence.

nammn · May 19, 2016 · ee1a999 · ee1a999
1 parent 59caabc
commit ee1a999
Show file tree

Hide file tree

Showing 8 changed files with 299 additions and 346 deletions.
diff --git a/component/all/payload.go b/component/all/payload.go
@@ -201,14 +201,16 @@ func (payloads) registerState() {
 		return
 	}
 
-	newUnitPayloads := func(persist state.Persistence, unit, machine string) (state.UnitPayloads, error) {
-		return payloadstate.NewUnitPayloads(persist, unit, machine), nil
+	newUnitPayloads := func(db state.Persistence, unit, machine string) (state.UnitPayloads, error) {
+		persist := persistence.NewPersistence(db)
+		unitPersist := persistence.NewUnitPersistence(persist, unit)
+		return payloadstate.NewUnitPayloads(unitPersist, unit, machine), nil
 	}
 
 	newEnvPayloads := func(db state.Persistence) (state.EnvPayloads, error) {
-		envPersist := persistence.NewEnvPersistence(db)
+		persist := persistence.NewPersistence(db)
 		envPayloads := payloadstate.EnvPayloads{
-			Persist: envPersist,
+			Persist: persist,
 		}
 		return envPayloads, nil
 	}

diff --git a/payload/persistence/base_test.go b/payload/persistence/base_test.go
@@ -47,8 +47,9 @@ func (s *BaseSuite) NewDoc(id string, pl payload.FullPayloadInfo) *payloadDoc {
 		MachineID: pl.Machine,
 		StateID:   id,
 		Type:      pl.Type,
-		RawID:     pl.ID,
 		State:     pl.Status,
+		Labels:    append([]string{}, pl.Labels...),
+		RawID:     pl.ID,
 	}
 }
 
@@ -64,7 +65,7 @@ func (s *BaseSuite) RemoveDoc(name string) {
 }
 
 func (s *BaseSuite) NewPersistence() *Persistence {
-	return NewPersistence(s.State, s.Unit)
+	return NewPersistence(s.State)
 }
 
 func (s *BaseSuite) SetUnit(id string) {

diff --git a/payload/persistence/env.go b/payload/persistence/env.go
diff --git a/payload/persistence/env_test.go b/payload/persistence/env_test.go