Modify unit tests for state to run on all state implementations

Signed-off-by: Matthew Heon <matthew.heon@gmail.com>

Closes: #229
Approved by: rhatdan

1 files changed, 24 insertions, 0 deletions

diff --git a/libpod/in_memory_state.go b/libpod/in_memory_state.go
index 244a1ab25..19d14366c 100644
--- a/libpod/in_memory_state.go
+++ b/libpod/in_memory_state.go
@@ -182,6 +182,18 @@ func (s *InMemoryState) RemoveContainer(ctr *Container) error {
 // As all state is in-memory, no update will be required
 // As such this is a no-op
 func (s *InMemoryState) UpdateContainer(ctr *Container) error {
+	// If the container is invalid, return error
+	if !ctr.valid {
+		return errors.Wrapf(ErrCtrRemoved, "container with ID %s is not valid", ctr.ID())
+	}
+
+	// If the container does not exist, return error
+	_, ok := s.containers[ctr.ID()]
+	if !ok {
+		ctr.valid = false
+		return errors.Wrapf(ErrNoSuchCtr, "container with ID %s not found in state", ctr.ID())
+	}
+
 	return nil
 }
 
@@ -190,6 +202,18 @@ func (s *InMemoryState) UpdateContainer(ctr *Container) error {
 // are made
 // As such this is a no-op
 func (s *InMemoryState) SaveContainer(ctr *Container) error {
+	// If the container is invalid, return error
+	if !ctr.valid {
+		return errors.Wrapf(ErrCtrRemoved, "container with ID %s is not valid", ctr.ID())
+	}
+
+	// If the container does not exist, return error
+	_, ok := s.containers[ctr.ID()]
+	if !ok {
+		ctr.valid = false
+		return errors.Wrapf(ErrNoSuchCtr, "container with ID %s not found in state", ctr.ID())
+	}
+
 	return nil
 }