Refactor storage to support shallow types.

Add Client to storage.
Fix client resizing issues.

storage/file.go

@@ -2,9 +2,7 @@ package storage
 
 import (
 	"encoding/json"
-	"fmt"
 	"os"
-	"strings"
 	"sync"
 
 	"github.com/play-with-docker/play-with-docker/pwd/types"
@@ -13,146 +11,76 @@ import (
 type storage struct {
 	rw   sync.Mutex
 	path string
-	db   map[string]*types.Session
+	db   *DB
 }
 
-func (store *storage) SessionGet(sessionId string) (*types.Session, error) {
+type DB struct {
+	Sessions         map[string]*types.Session         `json:"sessions"`
+	Instances        map[string]*types.Instance        `json:"instances"`
+	Clients          map[string]*types.Client          `json:"clients"`
+	WindowsInstances map[string]*types.WindowsInstance `json:"windows_instances"`
+
+	WindowsInstancesBySessionId map[string][]string `json:"windows_instances_by_session_id"`
+	InstancesBySessionId        map[string][]string `json:"instances_by_session_id"`
+	ClientsBySessionId          map[string][]string `json:"clients_by_session_id"`
+}
+
+func (store *storage) SessionGet(id string) (*types.Session, error) {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	s, found := store.db[sessionId]
+	s, found := store.db.Sessions[id]
 	if !found {
-		return nil, fmt.Errorf("%s", notFound)
+		return nil, notFound
 	}
 
 	return s, nil
 }
 
-func (store *storage) SessionGetAll() (map[string]*types.Session, error) {
+func (store *storage) SessionGetAll() ([]*types.Session, error) {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	return store.db, nil
+	sessions := make([]*types.Session, len(store.db.Sessions))
+	i := 0
+	for _, s := range store.db.Sessions {
+		sessions[i] = s
+		i++
+	}
+
+	return sessions, nil
 }
 
-func (store *storage) SessionPut(s *types.Session) error {
+func (store *storage) SessionPut(session *types.Session) error {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	// Initialize instances map if nil
-	if s.Instances == nil {
-		s.Instances = map[string]*types.Instance{}
-	}
-	store.db[s.Id] = s
+	store.db.Sessions[session.Id] = session
 
 	return store.save()
 }
 
-func (store *storage) InstanceGetAllWindows() ([]*types.WindowsInstance, error) {
+func (store *storage) SessionDelete(id string) error {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	instances := []*types.WindowsInstance{}
-
-	for _, s := range store.db {
-		instances = append(instances, s.WindowsAssigned...)
-	}
-
-	return instances, nil
-}
-
-func (store *storage) InstanceGet(sessionId, name string) (*types.Instance, error) {
-	store.rw.Lock()
-	defer store.rw.Unlock()
-
-	s := store.db[sessionId]
-	if s == nil {
-		return nil, fmt.Errorf("%s", notFound)
-	}
-	i := s.Instances[name]
-	if i == nil {
-		return nil, fmt.Errorf("%s", notFound)
-	}
-	return i, nil
-}
-
-func (store *storage) InstanceFindByIP(sessionId, ip string) (*types.Instance, error) {
-	store.rw.Lock()
-	defer store.rw.Unlock()
-
-	for id, s := range store.db {
-		if strings.HasPrefix(id, sessionId[:8]) {
-			for _, i := range s.Instances {
-				if i.IP == ip {
-					return i, nil
-				}
-			}
-		}
-	}
-
-	return nil, fmt.Errorf("%s", notFound)
-}
-
-func (store *storage) InstanceCreate(sessionId string, instance *types.Instance) error {
-	store.rw.Lock()
-	defer store.rw.Unlock()
-
-	s, found := store.db[sessionId]
+	_, found := store.db.Sessions[id]
 	if !found {
-		return fmt.Errorf("Session %s", notFound)
-	}
-
-	s.Instances[instance.Name] = instance
-
-	return store.save()
-}
-
-func (store *storage) InstanceCreateWindows(instance *types.WindowsInstance) error {
-	store.rw.Lock()
-	defer store.rw.Unlock()
-
-	s, found := store.db[instance.SessionId]
-	if !found {
-		return fmt.Errorf("Session %s", notFound)
-	}
-
-	s.WindowsAssigned = append(s.WindowsAssigned, instance)
-
-	return store.save()
-}
-
-func (store *storage) InstanceDelete(sessionId, name string) error {
-	store.rw.Lock()
-	defer store.rw.Unlock()
-
-	s, found := store.db[sessionId]
-	if !found {
-		return fmt.Errorf("Session %s", notFound)
-	}
-
-	if _, found := s.Instances[name]; !found {
 		return nil
 	}
-	delete(s.Instances, name)
-
-	return store.save()
-}
-
-func (store *storage) InstanceDeleteWindows(sessionId, id string) error {
-	store.rw.Lock()
-	defer store.rw.Unlock()
-
-	s, found := store.db[sessionId]
-	if !found {
-		return fmt.Errorf("Session %s", notFound)
+	for _, i := range store.db.WindowsInstancesBySessionId[id] {
+		delete(store.db.WindowsInstances, i)
 	}
-
-	for i, winst := range s.WindowsAssigned {
-		if winst.ID == id {
-			s.WindowsAssigned = append(s.WindowsAssigned[:i], s.WindowsAssigned[i+1:]...)
-		}
-
+	store.db.WindowsInstancesBySessionId[id] = []string{}
+	for _, i := range store.db.InstancesBySessionId[id] {
+		delete(store.db.Instances, i)
 	}
+	store.db.InstancesBySessionId[id] = []string{}
+	for _, i := range store.db.ClientsBySessionId[id] {
+		delete(store.db.Clients, i)
+	}
+	store.db.ClientsBySessionId[id] = []string{}
+	delete(store.db.Sessions, id)
 
 	return store.save()
 }
@@ -161,30 +89,217 @@ func (store *storage) SessionCount() (int, error) {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	return len(store.db), nil
+	return len(store.db.Sessions), nil
+}
+
+func (store *storage) InstanceGet(name string) (*types.Instance, error) {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	i := store.db.Instances[name]
+	if i == nil {
+		return nil, notFound
+	}
+	return i, nil
+}
+
+func (store *storage) InstancePut(instance *types.Instance) error {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	_, found := store.db.Sessions[string(instance.SessionId)]
+	if !found {
+		return notFound
+	}
+
+	store.db.Instances[instance.Name] = instance
+	found = false
+	for _, i := range store.db.InstancesBySessionId[string(instance.SessionId)] {
+		if i == instance.Name {
+			found = true
+			break
+		}
+	}
+	if !found {
+		store.db.InstancesBySessionId[string(instance.SessionId)] = append(store.db.InstancesBySessionId[string(instance.SessionId)], instance.Name)
+	}
+
+	return store.save()
+}
+
+func (store *storage) InstanceDelete(name string) error {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	instance, found := store.db.Instances[name]
+	if !found {
+		return nil
+	}
+
+	instances := store.db.InstancesBySessionId[string(instance.SessionId)]
+	for n, i := range instances {
+		if i == name {
+			instances = append(instances[:n], instances[n+1:]...)
+			break
+		}
+	}
+	store.db.InstancesBySessionId[string(instance.SessionId)] = instances
+	delete(store.db.Instances, name)
+
+	return store.save()
 }
 
 func (store *storage) InstanceCount() (int, error) {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	var ins int
-
-	for _, s := range store.db {
-		ins += len(s.Instances)
-	}
-
-	return ins, nil
+	return len(store.db.Instances), nil
 }
 
-func (store *storage) SessionDelete(sessionId string) error {
+func (store *storage) InstanceFindBySessionId(sessionId string) ([]*types.Instance, error) {
 	store.rw.Lock()
 	defer store.rw.Unlock()
 
-	delete(store.db, sessionId)
+	instanceIds := store.db.InstancesBySessionId[sessionId]
+	instances := make([]*types.Instance, len(instanceIds))
+	for i, id := range instanceIds {
+		instances[i] = store.db.Instances[id]
+	}
+
+	return instances, nil
+}
+
+func (store *storage) WindowsInstanceGetAll() ([]*types.WindowsInstance, error) {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	instances := []*types.WindowsInstance{}
+
+	for _, s := range store.db.WindowsInstances {
+		instances = append(instances, s)
+	}
+
+	return instances, nil
+}
+
+func (store *storage) WindowsInstancePut(instance *types.WindowsInstance) error {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	_, found := store.db.Sessions[string(instance.SessionId)]
+	if !found {
+		return notFound
+	}
+	store.db.WindowsInstances[instance.Id] = instance
+	found = false
+	for _, i := range store.db.WindowsInstancesBySessionId[string(instance.SessionId)] {
+		if i == instance.Id {
+			found = true
+			break
+		}
+	}
+	if !found {
+		store.db.WindowsInstancesBySessionId[string(instance.SessionId)] = append(store.db.WindowsInstancesBySessionId[string(instance.SessionId)], instance.Id)
+	}
+
 	return store.save()
 }
 
+func (store *storage) WindowsInstanceDelete(id string) error {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	instance, found := store.db.WindowsInstances[id]
+	if !found {
+		return nil
+	}
+
+	instances := store.db.WindowsInstancesBySessionId[string(instance.SessionId)]
+	for n, i := range instances {
+		if i == id {
+			instances = append(instances[:n], instances[n+1:]...)
+			break
+		}
+	}
+	store.db.WindowsInstancesBySessionId[string(instance.SessionId)] = instances
+	delete(store.db.WindowsInstances, id)
+
+	return store.save()
+}
+
+func (store *storage) ClientGet(id string) (*types.Client, error) {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	i := store.db.Clients[id]
+	if i == nil {
+		return nil, notFound
+	}
+	return i, nil
+}
+func (store *storage) ClientPut(client *types.Client) error {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	_, found := store.db.Sessions[string(client.SessionId)]
+	if !found {
+		return notFound
+	}
+
+	store.db.Clients[client.Id] = client
+	found = false
+	for _, i := range store.db.ClientsBySessionId[string(client.SessionId)] {
+		if i == client.Id {
+			found = true
+			break
+		}
+	}
+	if !found {
+		store.db.ClientsBySessionId[string(client.SessionId)] = append(store.db.ClientsBySessionId[string(client.SessionId)], client.Id)
+	}
+
+	return store.save()
+}
+func (store *storage) ClientDelete(id string) error {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	client, found := store.db.Clients[id]
+	if !found {
+		return nil
+	}
+
+	clients := store.db.ClientsBySessionId[string(client.SessionId)]
+	for n, i := range clients {
+		if i == client.Id {
+			clients = append(clients[:n], clients[n+1:]...)
+			break
+		}
+	}
+	store.db.ClientsBySessionId[string(client.SessionId)] = clients
+	delete(store.db.Clients, id)
+
+	return store.save()
+}
+func (store *storage) ClientCount() (int, error) {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	return len(store.db.Clients), nil
+}
+func (store *storage) ClientFindBySessionId(sessionId string) ([]*types.Client, error) {
+	store.rw.Lock()
+	defer store.rw.Unlock()
+
+	clientIds := store.db.ClientsBySessionId[sessionId]
+	clients := make([]*types.Client, len(clientIds))
+	for i, id := range clientIds {
+		clients[i] = store.db.Clients[id]
+	}
+
+	return clients, nil
+}
+
 func (store *storage) load() error {
 	file, err := os.Open(store.path)
 
@@ -196,7 +311,15 @@ func (store *storage) load() error {
 			return err
 		}
 	} else {
-		store.db = map[string]*types.Session{}
+		store.db = &DB{
+			Sessions:                    map[string]*types.Session{},
+			Instances:                   map[string]*types.Instance{},
+			Clients:                     map[string]*types.Client{},
+			WindowsInstances:            map[string]*types.WindowsInstance{},
+			WindowsInstancesBySessionId: map[string][]string{},
+			InstancesBySessionId:        map[string][]string{},
+			ClientsBySessionId:          map[string][]string{},
+		}
 	}
 
 	file.Close()

storage/file_test.go

@@ -29,9 +29,16 @@ func TestSessionPut(t *testing.T) {
 
 	assert.Nil(t, err)
 
-	var loadedSessions map[string]*types.Session
-	expectedSessions := map[string]*types.Session{}
-	expectedSessions[s.Id] = s
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{s.Id: s},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{},
+	}
+	var loadedDB *DB
 
 	file, err := os.Open(tmpfile.Name())
 
@@ -39,24 +46,31 @@ func TestSessionPut(t *testing.T) {
 	defer file.Close()
 
 	decoder := json.NewDecoder(file)
-	err = decoder.Decode(&loadedSessions)
+	err = decoder.Decode(&loadedDB)
 
 	assert.Nil(t, err)
 
-	assert.EqualValues(t, expectedSessions, loadedSessions)
+	assert.EqualValues(t, expectedDB, loadedDB)
 }
 
 func TestSessionGet(t *testing.T) {
-	expectedSession := &types.Session{Id: "session1"}
-	sessions := map[string]*types.Session{}
-	sessions[expectedSession.Id] = expectedSession
+	expectedSession := &types.Session{Id: "aaabbbccc"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{expectedSession.Id: expectedSession},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{},
+	}
 
 	tmpfile, err := ioutil.TempFile("", "pwd")
 	if err != nil {
 		log.Fatal(err)
 	}
 	encoder := json.NewEncoder(tmpfile)
-	err = encoder.Encode(&sessions)
+	err = encoder.Encode(&expectedDB)
 	assert.Nil(t, err)
 	tmpfile.Close()
 	defer os.Remove(tmpfile.Name())
@@ -65,28 +79,34 @@ func TestSessionGet(t *testing.T) {
 
 	assert.Nil(t, err)
 
-	_, err = storage.SessionGet("bad id")
+	_, err = storage.SessionGet("foobar")
 	assert.True(t, NotFound(err))
 
-	loadedSession, err := storage.SessionGet("session1")
+	loadedSession, err := storage.SessionGet("aaabbbccc")
 	assert.Nil(t, err)
 
 	assert.Equal(t, expectedSession, loadedSession)
 }
 
 func TestSessionGetAll(t *testing.T) {
-	s1 := &types.Session{Id: "session1"}
-	s2 := &types.Session{Id: "session2"}
-	sessions := map[string]*types.Session{}
-	sessions[s1.Id] = s1
-	sessions[s2.Id] = s2
+	s1 := &types.Session{Id: "aaabbbccc"}
+	s2 := &types.Session{Id: "dddeeefff"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{s1.Id: s1, s2.Id: s2},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{},
+	}
 
 	tmpfile, err := ioutil.TempFile("", "pwd")
 	if err != nil {
 		log.Fatal(err)
 	}
 	encoder := json.NewEncoder(tmpfile)
-	err = encoder.Encode(&sessions)
+	err = encoder.Encode(&expectedDB)
 	assert.Nil(t, err)
 	tmpfile.Close()
 	defer os.Remove(tmpfile.Name())
@@ -95,216 +115,11 @@ func TestSessionGetAll(t *testing.T) {
 
 	assert.Nil(t, err)
 
-	loadedSessions, err := storage.SessionGetAll()
+	sessions, err := storage.SessionGetAll()
 	assert.Nil(t, err)
 
-	assert.Equal(t, s1, loadedSessions[s1.Id])
-	assert.Equal(t, s2, loadedSessions[s2.Id])
-}
-
-func TestInstanceFindByIP(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-
-	i1 := &types.Instance{Name: "i1", IP: "10.0.0.1"}
-	i2 := &types.Instance{Name: "i2", IP: "10.1.0.1"}
-	s1 := &types.Session{Id: "session1", Instances: map[string]*types.Instance{"i1": i1}}
-	s2 := &types.Session{Id: "session2", Instances: map[string]*types.Instance{"i2": i2}}
-	err = storage.SessionPut(s1)
-	assert.Nil(t, err)
-	err = storage.SessionPut(s2)
-	assert.Nil(t, err)
-
-	foundInstance, err := storage.InstanceFindByIP("session1", "10.0.0.1")
-	assert.Nil(t, err)
-	assert.Equal(t, i1, foundInstance)
-
-	foundInstance, err = storage.InstanceFindByIP("session2", "10.1.0.1")
-	assert.Nil(t, err)
-	assert.Equal(t, i2, foundInstance)
-
-	foundInstance, err = storage.InstanceFindByIP("session3", "10.1.0.1")
-	assert.True(t, NotFound(err))
-	assert.Nil(t, foundInstance)
-
-	foundInstance, err = storage.InstanceFindByIP("session1", "10.1.0.1")
-	assert.True(t, NotFound(err))
-	assert.Nil(t, foundInstance)
-
-	foundInstance, err = storage.InstanceFindByIP("session1", "192.168.0.1")
-	assert.True(t, NotFound(err))
-	assert.Nil(t, foundInstance)
-}
-
-func TestInstanceGet(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-
-	i1 := &types.Instance{Name: "i1", IP: "10.0.0.1"}
-	s1 := &types.Session{Id: "session1", Instances: map[string]*types.Instance{"i1": i1}}
-	err = storage.SessionPut(s1)
-	assert.Nil(t, err)
-
-	foundInstance, err := storage.InstanceGet("session1", "i1")
-	assert.Nil(t, err)
-	assert.Equal(t, i1, foundInstance)
-}
-
-func TestInstanceGetAllWindows(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-	w1 := []*types.WindowsInstance{{ID: "one"}, {ID: "two"}}
-	w2 := []*types.WindowsInstance{{ID: "three"}, {ID: "four"}}
-	s1 := &types.Session{Id: "session1", WindowsAssigned: w1}
-	s2 := &types.Session{Id: "session2", WindowsAssigned: w2}
-	err = storage.SessionPut(s1)
-	err = storage.SessionPut(s2)
-	assert.Nil(t, err)
-
-	allw, err := storage.InstanceGetAllWindows()
-	assert.Nil(t, err)
-	assert.Equal(t, allw, append(w1, w2...))
-}
-
-func TestInstanceCreate(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-
-	i1 := &types.Instance{Name: "i1", IP: "10.0.0.1"}
-	s1 := &types.Session{Id: "session1"}
-	err = storage.SessionPut(s1)
-	assert.Nil(t, err)
-	err = storage.InstanceCreate(s1.Id, i1)
-	assert.Nil(t, err)
-
-	loadedSession, err := storage.SessionGet("session1")
-	assert.Nil(t, err)
-
-	assert.Equal(t, i1, loadedSession.Instances["i1"])
-
-}
-
-func TestInstanceCreateWindows(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-
-	s1 := &types.Session{Id: "session1"}
-	i1 := &types.WindowsInstance{SessionId: s1.Id, ID: "some id"}
-	err = storage.SessionPut(s1)
-	assert.Nil(t, err)
-	err = storage.InstanceCreateWindows(i1)
-	assert.Nil(t, err)
-
-	loadedSession, err := storage.SessionGet("session1")
-	assert.Nil(t, err)
-
-	assert.Equal(t, i1, loadedSession.WindowsAssigned[0])
-}
-
-func TestInstanceDeleteWindows(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-
-	s1 := &types.Session{Id: "session1", WindowsAssigned: []*types.WindowsInstance{{ID: "one"}}}
-	err = storage.SessionPut(s1)
-	assert.Nil(t, err)
-
-	err = storage.InstanceDeleteWindows(s1.Id, "one")
-	assert.Nil(t, err)
-
-	found, err := storage.SessionGet(s1.Id)
-	assert.Equal(t, 0, len(found.WindowsAssigned))
-}
-
-func TestCounts(t *testing.T) {
-	tmpfile, err := ioutil.TempFile("", "pwd")
-	if err != nil {
-		log.Fatal(err)
-	}
-	tmpfile.Close()
-	os.Remove(tmpfile.Name())
-	defer os.Remove(tmpfile.Name())
-
-	storage, err := NewFileStorage(tmpfile.Name())
-
-	assert.Nil(t, err)
-
-	c1 := &types.Client{}
-	i1 := &types.Instance{Name: "i1", IP: "10.0.0.1"}
-	i2 := &types.Instance{Name: "i2", IP: "10.1.0.1"}
-	s1 := &types.Session{Id: "session1", Instances: map[string]*types.Instance{"i1": i1}}
-	s2 := &types.Session{Id: "session2", Instances: map[string]*types.Instance{"i2": i2}}
-	s3 := &types.Session{Id: "session3", Clients: []*types.Client{c1}}
-
-	err = storage.SessionPut(s1)
-	assert.Nil(t, err)
-	err = storage.SessionPut(s2)
-	assert.Nil(t, err)
-	err = storage.SessionPut(s3)
-	assert.Nil(t, err)
-
-	num, err := storage.SessionCount()
-	assert.Nil(t, err)
-	assert.Equal(t, 3, num)
-
-	num, err = storage.InstanceCount()
-	assert.Nil(t, err)
-	assert.Equal(t, 2, num)
-
+	assert.Subset(t, sessions, []*types.Session{s1, s2})
+	assert.Len(t, sessions, 2)
 }
 
 func TestSessionDelete(t *testing.T) {
@@ -335,3 +150,401 @@ func TestSessionDelete(t *testing.T) {
 	assert.True(t, NotFound(err))
 	assert.Nil(t, found)
 }
+
+func TestInstanceGet(t *testing.T) {
+	expectedInstance := &types.Instance{SessionId: "aaabbbccc", Name: "i1", IP: "10.0.0.1"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{},
+		Instances:                   map[string]*types.Instance{expectedInstance.Name: expectedInstance},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{expectedInstance.SessionId: []string{expectedInstance.Name}},
+		ClientsBySessionId:          map[string][]string{},
+	}
+
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	encoder := json.NewEncoder(tmpfile)
+	err = encoder.Encode(&expectedDB)
+	assert.Nil(t, err)
+	tmpfile.Close()
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	foundInstance, err := storage.InstanceGet("i1")
+	assert.Nil(t, err)
+	assert.Equal(t, expectedInstance, foundInstance)
+}
+
+func TestInstancePut(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	tmpfile.Close()
+	os.Remove(tmpfile.Name())
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	s := &types.Session{Id: "aaabbbccc"}
+	i := &types.Instance{Name: "i1", IP: "10.0.0.1", SessionId: s.Id}
+
+	err = storage.SessionPut(s)
+	assert.Nil(t, err)
+
+	err = storage.InstancePut(i)
+	assert.Nil(t, err)
+
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{s.Id: s},
+		Instances:                   map[string]*types.Instance{i.Name: i},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{i.SessionId: []string{i.Name}},
+		ClientsBySessionId:          map[string][]string{},
+	}
+	var loadedDB *DB
+
+	file, err := os.Open(tmpfile.Name())
+
+	assert.Nil(t, err)
+	defer file.Close()
+
+	decoder := json.NewDecoder(file)
+	err = decoder.Decode(&loadedDB)
+
+	assert.Nil(t, err)
+
+	assert.EqualValues(t, expectedDB, loadedDB)
+}
+
+func TestInstanceDelete(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	tmpfile.Close()
+	os.Remove(tmpfile.Name())
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	s := &types.Session{Id: "session1"}
+	err = storage.SessionPut(s)
+	assert.Nil(t, err)
+
+	i := &types.Instance{Name: "i1", IP: "10.0.0.1", SessionId: s.Id}
+	err = storage.InstancePut(i)
+	assert.Nil(t, err)
+
+	found, err := storage.InstanceGet(i.Name)
+	assert.Nil(t, err)
+	assert.Equal(t, i, found)
+
+	err = storage.InstanceDelete(i.Name)
+	assert.Nil(t, err)
+
+	found, err = storage.InstanceGet(i.Name)
+	assert.True(t, NotFound(err))
+	assert.Nil(t, found)
+}
+
+func TestInstanceFindBySessionId(t *testing.T) {
+	i1 := &types.Instance{SessionId: "aaabbbccc", Name: "c1"}
+	i2 := &types.Instance{SessionId: "aaabbbccc", Name: "c2"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{},
+		Instances:                   map[string]*types.Instance{i1.Name: i1, i2.Name: i2},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{i1.SessionId: []string{i1.Name, i2.Name}},
+		ClientsBySessionId:          map[string][]string{},
+	}
+
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	encoder := json.NewEncoder(tmpfile)
+	err = encoder.Encode(&expectedDB)
+	assert.Nil(t, err)
+	tmpfile.Close()
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	instances, err := storage.InstanceFindBySessionId("aaabbbccc")
+	assert.Nil(t, err)
+	assert.Subset(t, instances, []*types.Instance{i1, i2})
+	assert.Len(t, instances, 2)
+}
+
+func TestWindowsInstanceGetAll(t *testing.T) {
+	i1 := &types.WindowsInstance{SessionId: "aaabbbccc", Id: "i1"}
+	i2 := &types.WindowsInstance{SessionId: "aaabbbccc", Id: "i2"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{i1.Id: i1, i2.Id: i2},
+		WindowsInstancesBySessionId: map[string][]string{i1.SessionId: []string{i1.Id, i2.Id}},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{},
+	}
+
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	encoder := json.NewEncoder(tmpfile)
+	err = encoder.Encode(&expectedDB)
+	assert.Nil(t, err)
+	tmpfile.Close()
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	instances, err := storage.WindowsInstanceGetAll()
+	assert.Nil(t, err)
+	assert.Subset(t, instances, []*types.WindowsInstance{i1, i2})
+	assert.Len(t, instances, 2)
+}
+
+func TestWindowsInstancePut(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	tmpfile.Close()
+	os.Remove(tmpfile.Name())
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	s := &types.Session{Id: "aaabbbccc"}
+	i := &types.WindowsInstance{Id: "i1", SessionId: s.Id}
+
+	err = storage.SessionPut(s)
+	assert.Nil(t, err)
+
+	err = storage.WindowsInstancePut(i)
+	assert.Nil(t, err)
+
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{s.Id: s},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{},
+		WindowsInstances:            map[string]*types.WindowsInstance{i.Id: i},
+		WindowsInstancesBySessionId: map[string][]string{i.SessionId: []string{i.Id}},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{},
+	}
+	var loadedDB *DB
+
+	file, err := os.Open(tmpfile.Name())
+
+	assert.Nil(t, err)
+	defer file.Close()
+
+	decoder := json.NewDecoder(file)
+	err = decoder.Decode(&loadedDB)
+
+	assert.Nil(t, err)
+
+	assert.EqualValues(t, expectedDB, loadedDB)
+}
+
+func TestWindowsInstanceDelete(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	tmpfile.Close()
+	os.Remove(tmpfile.Name())
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	s := &types.Session{Id: "session1"}
+	err = storage.SessionPut(s)
+	assert.Nil(t, err)
+
+	i := &types.WindowsInstance{Id: "i1", SessionId: s.Id}
+	err = storage.WindowsInstancePut(i)
+	assert.Nil(t, err)
+
+	found, err := storage.WindowsInstanceGetAll()
+	assert.Nil(t, err)
+	assert.Equal(t, []*types.WindowsInstance{i}, found)
+
+	err = storage.WindowsInstanceDelete(i.Id)
+	assert.Nil(t, err)
+
+	found, err = storage.WindowsInstanceGetAll()
+	assert.Nil(t, err)
+	assert.Empty(t, found)
+}
+
+func TestClientGet(t *testing.T) {
+	c := &types.Client{SessionId: "aaabbbccc", Id: "c1"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{c.Id: c},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{c.SessionId: []string{c.Id}},
+	}
+
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	encoder := json.NewEncoder(tmpfile)
+	err = encoder.Encode(&expectedDB)
+	assert.Nil(t, err)
+	tmpfile.Close()
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	found, err := storage.ClientGet("c1")
+	assert.Nil(t, err)
+	assert.Equal(t, c, found)
+}
+
+func TestClientPut(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	tmpfile.Close()
+	os.Remove(tmpfile.Name())
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	s := &types.Session{Id: "aaabbbccc"}
+	c := &types.Client{Id: "c1", SessionId: s.Id}
+
+	err = storage.SessionPut(s)
+	assert.Nil(t, err)
+
+	err = storage.ClientPut(c)
+	assert.Nil(t, err)
+
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{s.Id: s},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{c.Id: c},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{c.SessionId: []string{c.Id}},
+	}
+	var loadedDB *DB
+
+	file, err := os.Open(tmpfile.Name())
+
+	assert.Nil(t, err)
+	defer file.Close()
+
+	decoder := json.NewDecoder(file)
+	err = decoder.Decode(&loadedDB)
+
+	assert.Nil(t, err)
+
+	assert.EqualValues(t, expectedDB, loadedDB)
+}
+
+func TestClientDelete(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	tmpfile.Close()
+	os.Remove(tmpfile.Name())
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	s := &types.Session{Id: "session1"}
+	err = storage.SessionPut(s)
+	assert.Nil(t, err)
+
+	c := &types.Client{Id: "c1", SessionId: s.Id}
+	err = storage.ClientPut(c)
+	assert.Nil(t, err)
+
+	found, err := storage.ClientGet(c.Id)
+	assert.Nil(t, err)
+	assert.Equal(t, c, found)
+
+	err = storage.ClientDelete(c.Id)
+	assert.Nil(t, err)
+
+	found, err = storage.ClientGet(c.Id)
+	assert.True(t, NotFound(err))
+	assert.Nil(t, found)
+}
+
+func TestClientFindBySessionId(t *testing.T) {
+	c1 := &types.Client{SessionId: "aaabbbccc", Id: "c1"}
+	c2 := &types.Client{SessionId: "aaabbbccc", Id: "c2"}
+	expectedDB := &DB{
+		Sessions:                    map[string]*types.Session{},
+		Instances:                   map[string]*types.Instance{},
+		Clients:                     map[string]*types.Client{c1.Id: c1, c2.Id: c2},
+		WindowsInstances:            map[string]*types.WindowsInstance{},
+		WindowsInstancesBySessionId: map[string][]string{},
+		InstancesBySessionId:        map[string][]string{},
+		ClientsBySessionId:          map[string][]string{c1.SessionId: []string{c1.Id, c2.Id}},
+	}
+
+	tmpfile, err := ioutil.TempFile("", "pwd")
+	if err != nil {
+		log.Fatal(err)
+	}
+	encoder := json.NewEncoder(tmpfile)
+	err = encoder.Encode(&expectedDB)
+	assert.Nil(t, err)
+	tmpfile.Close()
+	defer os.Remove(tmpfile.Name())
+
+	storage, err := NewFileStorage(tmpfile.Name())
+
+	assert.Nil(t, err)
+
+	clients, err := storage.ClientFindBySessionId("aaabbbccc")
+	assert.Nil(t, err)
+	assert.Subset(t, clients, []*types.Client{c1, c2})
+	assert.Len(t, clients, 2)
+}

storage/mock.go

@@ -9,67 +9,76 @@ type Mock struct {
 	mock.Mock
 }
 
-func (m *Mock) SessionGet(sessionId string) (*types.Session, error) {
-	args := m.Called(sessionId)
+func (m *Mock) SessionGet(id string) (*types.Session, error) {
+	args := m.Called(id)
 	return args.Get(0).(*types.Session), args.Error(1)
 }
-
+func (m *Mock) SessionGetAll() ([]*types.Session, error) {
+	args := m.Called()
+	return args.Get(0).([]*types.Session), args.Error(1)
+}
 func (m *Mock) SessionPut(session *types.Session) error {
 	args := m.Called(session)
 	return args.Error(0)
 }
-
+func (m *Mock) SessionDelete(id string) error {
+	args := m.Called(id)
+	return args.Error(0)
+}
 func (m *Mock) SessionCount() (int, error) {
 	args := m.Called()
 	return args.Int(0), args.Error(1)
 }
-
-func (m *Mock) SessionDelete(sessionId string) error {
-	args := m.Called(sessionId)
-	return args.Error(0)
-}
-
-func (m *Mock) SessionGetAll() (map[string]*types.Session, error) {
-	args := m.Called()
-	return args.Get(0).(map[string]*types.Session), args.Error(1)
-}
-
-func (m *Mock) InstanceGet(sessionId, name string) (*types.Instance, error) {
-	args := m.Called(sessionId, name)
+func (m *Mock) InstanceGet(name string) (*types.Instance, error) {
+	args := m.Called(name)
 	return args.Get(0).(*types.Instance), args.Error(1)
 }
-
-func (m *Mock) InstanceGetAllWindows() ([]*types.WindowsInstance, error) {
-	args := m.Called()
-	return args.Get(0).([]*types.WindowsInstance), args.Error(1)
-}
-
-func (m *Mock) InstanceFindByIP(sessionId, ip string) (*types.Instance, error) {
-	args := m.Called(sessionId, ip)
-	return args.Get(0).(*types.Instance), args.Error(1)
-}
-
-func (m *Mock) InstanceCreate(sessionId string, instance *types.Instance) error {
-	args := m.Called(sessionId, instance)
-	return args.Error(0)
-}
-
-func (m *Mock) InstanceCreateWindows(instance *types.WindowsInstance) error {
+func (m *Mock) InstancePut(instance *types.Instance) error {
 	args := m.Called(instance)
 	return args.Error(0)
 }
-
-func (m *Mock) InstanceDelete(sessionId, instanceName string) error {
-	args := m.Called(sessionId, instanceName)
+func (m *Mock) InstanceDelete(name string) error {
+	args := m.Called(name)
 	return args.Error(0)
 }
-
-func (m *Mock) InstanceDeleteWindows(sessionId, instanceId string) error {
-	args := m.Called(sessionId, instanceId)
-	return args.Error(0)
-}
-
 func (m *Mock) InstanceCount() (int, error) {
 	args := m.Called()
 	return args.Int(0), args.Error(1)
 }
+func (m *Mock) InstanceFindBySessionId(sessionId string) ([]*types.Instance, error) {
+	args := m.Called(sessionId)
+	return args.Get(0).([]*types.Instance), args.Error(1)
+}
+
+func (m *Mock) WindowsInstanceGetAll() ([]*types.WindowsInstance, error) {
+	args := m.Called()
+	return args.Get(0).([]*types.WindowsInstance), args.Error(1)
+}
+func (m *Mock) WindowsInstancePut(instance *types.WindowsInstance) error {
+	args := m.Called(instance)
+	return args.Error(0)
+}
+func (m *Mock) WindowsInstanceDelete(id string) error {
+	args := m.Called(id)
+	return args.Error(0)
+}
+func (m *Mock) ClientGet(id string) (*types.Client, error) {
+	args := m.Called(id)
+	return args.Get(0).(*types.Client), args.Error(1)
+}
+func (m *Mock) ClientPut(client *types.Client) error {
+	args := m.Called(client)
+	return args.Error(0)
+}
+func (m *Mock) ClientDelete(id string) error {
+	args := m.Called(id)
+	return args.Error(0)
+}
+func (m *Mock) ClientCount() (int, error) {
+	args := m.Called()
+	return args.Int(0), args.Error(1)
+}
+func (m *Mock) ClientFindBySessionId(sessionId string) ([]*types.Client, error) {
+	args := m.Called(sessionId)
+	return args.Get(0).([]*types.Client), args.Error(1)
+}

storage/storage.go

@@ -1,34 +1,37 @@
 package storage
 
 import (
-	"fmt"
+	"errors"
 
 	"github.com/play-with-docker/play-with-docker/pwd/types"
 )
 
-const notFound = "NotFound"
+var notFound = errors.New("NotFound")
 
 func NotFound(e error) bool {
-	return e.Error() == notFound
-}
-
-func NewNotFoundError() error {
-	return fmt.Errorf("%s", notFound)
+	return e == notFound
 }
 
 type StorageApi interface {
-	SessionGet(string) (*types.Session, error)
-	SessionPut(*types.Session) error
+	SessionGet(id string) (*types.Session, error)
+	SessionGetAll() ([]*types.Session, error)
+	SessionPut(session *types.Session) error
+	SessionDelete(id string) error
 	SessionCount() (int, error)
-	SessionDelete(string) error
-	SessionGetAll() (map[string]*types.Session, error)
 
-	InstanceGet(sessionId, name string) (*types.Instance, error)
-	InstanceFindByIP(session, ip string) (*types.Instance, error)
-	InstanceCreate(sessionId string, instance *types.Instance) error
-	InstanceDelete(sessionId, instanceName string) error
-	InstanceDeleteWindows(sessionId, instanceId string) error
+	InstanceGet(name string) (*types.Instance, error)
+	InstancePut(instance *types.Instance) error
+	InstanceDelete(name string) error
 	InstanceCount() (int, error)
-	InstanceGetAllWindows() ([]*types.WindowsInstance, error)
-	InstanceCreateWindows(*types.WindowsInstance) error
+	InstanceFindBySessionId(sessionId string) ([]*types.Instance, error)
+
+	WindowsInstanceGetAll() ([]*types.WindowsInstance, error)
+	WindowsInstancePut(instance *types.WindowsInstance) error
+	WindowsInstanceDelete(id string) error
+
+	ClientGet(id string) (*types.Client, error)
+	ClientPut(client *types.Client) error
+	ClientDelete(id string) error
+	ClientCount() (int, error)
+	ClientFindBySessionId(sessionId string) ([]*types.Client, error)
 }