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

@@ -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
 }

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")
+	}
+}

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) {