From 70510b3f8688f50865be879c9d6775a2acaba0a7 Mon Sep 17 00:00:00 2001
From: Dongsu Park <dongsu@endocode.com>
Date: Fri, 21 Oct 2016 15:42:04 +0200
Subject: [PATCH] registry: add new unit tests for GetUnitState() and relevant
 ones

We should add new unit tests like TestUnitState() as well as
TestInMemoryUnitState(), to cover the new cases like
EtcdRegistry.UnitState() or inmemoryRegistry.UnitState().
---
 registry/fake.go              | 27 +++++++++++++
 registry/rpc/inmemory_test.go | 51 +++++++++++++++++++++++
 registry/unit_state_test.go   | 76 +++++++++++++++++++++++++++++++++++
 3 files changed, 154 insertions(+)

diff --git a/registry/fake.go b/registry/fake.go
index a78babd34..52be56eed 100644
--- a/registry/fake.go
+++ b/registry/fake.go
@@ -16,6 +16,7 @@ package registry
 
 import (
 	"errors"
+	"fmt"
 	"sort"
 	"sync"
 	"time"
@@ -281,6 +282,32 @@ func (f *FakeRegistry) UnitStates() ([]*unit.UnitState, error) {
 	return states, nil
 }
 
+func (f *FakeRegistry) UnitState(unitName string) (*unit.UnitState, error) {
+	f.Lock()
+	defer f.Unlock()
+
+	var us *unit.UnitState
+	for name := range f.jobStates {
+		sMIDs := make([]string, 0)
+		for machineID := range f.jobStates[name] {
+			sMIDs = append(sMIDs, machineID)
+		}
+		for _, mID := range sMIDs {
+			js := f.jobStates[name][mID]
+			if name == js.UnitName {
+				us = js
+				break
+			}
+		}
+	}
+
+	if us == nil {
+		return nil, fmt.Errorf("Cannot find unit(%s) from fakeRegistry", unitName)
+	}
+
+	return us, nil
+}
+
 func (f *FakeRegistry) UnitHeartbeat(name, machID string, ttl time.Duration) error {
 	return nil
 }
diff --git a/registry/rpc/inmemory_test.go b/registry/rpc/inmemory_test.go
index 36f5eda0f..87a816566 100644
--- a/registry/rpc/inmemory_test.go
+++ b/registry/rpc/inmemory_test.go
@@ -122,3 +122,54 @@ Description = Foo
 		t.Fatalf("unexpected amount of unit states in the in-memory registry got %d expected 1", len(inmemoryRegistry.UnitStates()))
 	}
 }
+
+func TestInMemoryUnitState(t *testing.T) {
+	inmemoryRegistry := newInmemoryRegistry()
+
+	scheduleUnit := &pb.ScheduledUnit{
+		Name:         "foo",
+		CurrentState: pb.TargetState_INACTIVE,
+		MachineID:    "machine1",
+	}
+	inmemoryRegistry.scheduledUnits[scheduleUnit.Name] = *scheduleUnit
+	contents := `
+[Unit]
+Description = Foo
+`
+	unitFile, err := unit.NewUnitFile(contents)
+	if err != nil {
+		t.Fatalf("unexpected error parsing unit %q: %v", contents, err)
+	}
+	unit := &pb.Unit{
+		Name:         "foo",
+		Unit:         unitFile.ToPB(),
+		DesiredState: pb.TargetState_LOADED,
+	}
+	machineID := "testMachine"
+	ttl := 2 * time.Second
+	inmemoryRegistry.CreateUnit(unit)
+	inmemoryRegistry.ScheduleUnit(unit.Name, machineID)
+
+	stateLoaded := &pb.UnitState{
+		Name:        unit.Name,
+		Hash:        "heh",
+		LoadState:   "active",
+		ActiveState: "loaded",
+		SubState:    "active",
+		MachineID:   machineID,
+	}
+
+	inmemoryRegistry.SaveUnitState(unit.Name, stateLoaded, ttl)
+	if !inmemoryRegistry.isUnitLoaded(unit.Name, machineID) {
+		u, ok := inmemoryRegistry.Unit(unit.Name)
+		if !ok {
+			t.Fatalf("unexpected error unit not found %s", unit.Name)
+		}
+		t.Fatalf("unexpected error unit expected to be loaded %v", u)
+	}
+
+	resUs := inmemoryRegistry.UnitState(unit.Name)
+	if resUs == nil || len(resUs.Name) == 0 {
+		t.Fatal("Invalid unit state in the in-memory registry")
+	}
+}
diff --git a/registry/unit_state_test.go b/registry/unit_state_test.go
index 3b5c4ba3d..46972624b 100644
--- a/registry/unit_state_test.go
+++ b/registry/unit_state_test.go
@@ -495,6 +495,82 @@ func TestUnitStates(t *testing.T) {
 	}
 }
 
+func TestUnitState(t *testing.T) {
+	fus1 := unit.UnitState{
+		LoadState:   "abc",
+		ActiveState: "def",
+		SubState:    "ghi",
+		MachineID:   "mID1",
+		UnitHash:    "zzz",
+		UnitName:    "foo",
+	}
+	// Multiple new unit states reported for the same unit
+	foo := &etcd.Node{
+		Key: "/fleet/states/foo",
+		Nodes: []*etcd.Node{
+			&etcd.Node{
+				Key:   "/fleet/states/foo/mID1",
+				Value: usToJson(t, &fus1),
+			},
+		},
+	}
+	// Bogus new unit state which we won't expect to see in results
+	bar := &etcd.Node{
+		Key: "/fleet/states/bar",
+		Nodes: []*etcd.Node{
+			&etcd.Node{
+				Key:   "/fleet/states/bar/asdf",
+				Value: `total garbage`,
+			},
+		},
+	}
+	// Response from crawling the new "states" namespace
+	res2 := &etcd.Response{
+		Node: &etcd.Node{
+			Key:   "/fleet/states",
+			Nodes: []*etcd.Node{foo, bar},
+		},
+	}
+	e := &testEtcdKeysAPI{
+		res: []*etcd.Response{res2},
+	}
+	r := &EtcdRegistry{kAPI: e, keyPrefix: "/fleet/"}
+
+	got, err := r.UnitState(fus1.UnitName)
+	if err != nil {
+		t.Errorf("unexpected error calling UnitState(%s): %v", fus1.UnitName, err)
+	}
+
+	want := &fus1
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("UnitState() returned unexpected result")
+		t.Log("got:")
+		t.Logf("%#v", got)
+		t.Log("want:")
+		t.Logf("%#v", want)
+	}
+
+	// Ensure UnitState handles different error scenarios appropriately
+	for i, tt := range []struct {
+		errs []error
+		fail bool
+	}{
+		{[]error{etcd.Error{Code: etcd.ErrorCodeKeyNotFound}}, false},
+		{[]error{etcd.Error{Code: etcd.ErrorCodeKeyNotFound}}, false},
+		{[]error{nil}, false}, // No errors, no responses should succeed
+		{[]error{errors.New("ur registry don't work")}, true},
+		{[]error{errors.New("ur registry don't work")}, true},
+	} {
+		e = &testEtcdKeysAPI{err: tt.errs}
+		r = &EtcdRegistry{kAPI: e, keyPrefix: "/fleet"}
+		got, err = r.UnitState(fus1.UnitName)
+		if (err != nil) != tt.fail {
+			t.Errorf("case %d: unexpected error state calling UnitState(%s): got %v, want %v",
+				i, fus1.UnitName, err, tt.fail)
+		}
+	}
+}
+
 func TestMUSKeys(t *testing.T) {
 	equal := func(a MUSKeys, b []MUSKey) bool {
 		if len(a) != len(b) {