Golang registry.NewFakeRegistry函數代碼示例

func TestMachinesList(t *testing.T) {
	fr := registry.NewFakeRegistry()
	fr.SetMachines([]machine.MachineState{
		{ID: "XXX", PublicIP: "", Metadata: nil},
		{ID: "YYY", PublicIP: "1.2.3.4", Metadata: map[string]string{"ping": "pong"}},
	})
	resource := &machinesResource{fr}
	rw := httptest.NewRecorder()
	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("Failed creating http.Request: %v", err)
	}

	resource.ServeHTTP(rw, req)
	if rw.Code != http.StatusOK {
		t.Errorf("Expected 200, got %d", rw.Code)
	}

	ct := rw.HeaderMap["Content-Type"]
	if len(ct) != 1 {
		t.Errorf("Response has wrong number of Content-Type values: %v", ct)
	} else if ct[0] != "application/json" {
		t.Errorf("Expected application/json, got %s", ct)
	}

	if rw.Body == nil {
		t.Error("Received nil response body")
	} else {
		body := rw.Body.String()
		expected := `{"machines":[{"id":"XXX"},{"id":"YYY","metadata":{"ping":"pong"},"primaryIP":"1.2.3.4"}]}`
		if body != expected {
			t.Errorf("Expected body:\n%s\n\nReceived body:\n%s\n", expected, body)
		}
	}
}

func newFakeRegistryForSsh() registry.Registry {
	machines := []machine.MachineState{
		{"c31e44e1-f858-436e-933e-59c642517860", "1.2.3.4", map[string]string{"ping": "pong"}, "", resource.ResourceTuple{}},
		{"595989bb-cbb7-49ce-8726-722d6e157b4e", "5.6.7.8", map[string]string{"foo": "bar"}, "", resource.ResourceTuple{}},
		{"hello.service", "8.7.6.5", map[string]string{"foo": "bar"}, "", resource.ResourceTuple{}},
	}

	jobs := []job.Job{
		*job.NewJob("j1.service", unit.Unit{}),
		*job.NewJob("j2.service", unit.Unit{}),
		*job.NewJob("hello.service", unit.Unit{}),
	}

	states := map[string]*unit.UnitState{
		"j1.service":    unit.NewUnitState("loaded", "active", "listening", &machines[0]),
		"j2.service":    unit.NewUnitState("loaded", "inactive", "dead", &machines[1]),
		"hello.service": unit.NewUnitState("loaded", "inactive", "dead", &machines[2]),
	}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(machines)
	reg.SetUnitStates(states)
	reg.SetJobs(jobs)

	return reg
}

func setupRegistryForStart(echoAttempts int) {
	m1 := machine.MachineState{
		ID:       "c31e44e1-f858-436e-933e-59c642517860",
		PublicIP: "1.2.3.4",
		Metadata: map[string]string{"ping": "pong"},
	}
	m2 := machine.MachineState{
		ID:       "595989bb-cbb7-49ce-8726-722d6e157b4e",
		PublicIP: "5.6.7.8",
		Metadata: map[string]string{"foo": "bar"},
	}
	m3 := machine.MachineState{
		ID:       "520983A8-FB9C-4A68-B49C-CED5BB2E9D08",
		Metadata: map[string]string{"foo": "bar"},
	}

	js := unit.NewUnitState("loaded", "active", "listening", m1.ID)
	js2 := unit.NewUnitState("loaded", "inactive", "dead", m2.ID)
	js3 := unit.NewUnitState("loaded", "inactive", "dead", m2.ID)
	js4 := unit.NewUnitState("loaded", "inactive", "dead", m3.ID)

	states := map[string]*unit.UnitState{"pong.service": js, "hello.service": js2, "echo.service": js3, "private.service": js4}
	machines := []machine.MachineState{m1, m2, m3}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(machines)
	reg.SetUnitStates(states)

	cAPI = &BlockedFakeRegistry{echoAttempts, *reg}
}

func TestDefaultHandlers(t *testing.T) {
	tests := []struct {
		method string
		path   string
		code   int
	}{
		{"GET", "/", http.StatusMethodNotAllowed},
		{"GET", "/v1-alpha", http.StatusMethodNotAllowed},
		{"GET", "/bogus", http.StatusNotFound},
	}

	for i, tt := range tests {
		fr := registry.NewFakeRegistry()
		hdlr := NewServeMux(fr)
		rr := httptest.NewRecorder()

		req, err := http.NewRequest(tt.method, tt.path, nil)
		if err != nil {
			t.Errorf("case %d: failed setting up http.Request for test: %v", i, err)
			continue
		}

		hdlr.ServeHTTP(rr, req)

		err = assertErrorResponse(rr, tt.code)
		if err != nil {
			t.Errorf("case %d: %v", i, err)
		}
	}
}

func newFakeRegistryForSsh() registry.Registry {
	// clear machineStates for every invocation
	machineStates = nil
	machines := []machine.MachineState{
		newMachineState("c31e44e1-f858-436e-933e-59c642517860", "1.2.3.4", map[string]string{"ping": "pong"}),
		newMachineState("595989bb-cbb7-49ce-8726-722d6e157b4e", "5.6.7.8", map[string]string{"foo": "bar"}),
		newMachineState("hello.service", "8.7.6.5", map[string]string{"foo": "bar"}),
	}

	jobs := []job.Job{
		*job.NewJob("j1.service", unit.Unit{}),
		*job.NewJob("j2.service", unit.Unit{}),
		*job.NewJob("hello.service", unit.Unit{}),
	}

	states := map[string]*unit.UnitState{
		"j1.service":    unit.NewUnitState("loaded", "active", "listening", machines[0].ID),
		"j2.service":    unit.NewUnitState("loaded", "inactive", "dead", machines[1].ID),
		"hello.service": unit.NewUnitState("loaded", "inactive", "dead", machines[2].ID),
	}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(machines)
	reg.SetUnitStates(states)
	reg.SetJobs(jobs)

	return reg
}

func TestAgentLoadUnloadJob(t *testing.T) {
	uManager := unit.NewFakeUnitManager()
	usGenerator := unit.NewUnitStateGenerator(uManager)
	fReg := registry.NewFakeRegistry()
	mach := &machine.FakeMachine{machine.MachineState{ID: "XXX"}}
	a, err := New(uManager, usGenerator, fReg, mach, DefaultTTL)
	if err != nil {
		t.Fatalf("Failed creating Agent: %v", err)
	}

	j := newTestJobFromUnitContents(t, "foo.service", "")
	err = a.loadJob(j)
	if err != nil {
		t.Fatalf("Failed calling Agent.loadJob: %v", err)
	}

	jobs, err := a.jobs()
	if err != nil {
		t.Fatalf("Failed calling Agent.jobs: %v", err)
	}

	jsLoaded := job.JobStateLoaded
	expectJobs := map[string]*job.Job{
		"foo.service": &job.Job{
			Name: "foo.service",
			UnitState: &unit.UnitState{
				LoadState:   "loaded",
				ActiveState: "active",
				SubState:    "running",
				MachineID:   "",
			},
			State: &jsLoaded,

			Unit:            unit.Unit{},
			TargetState:     job.JobState(""),
			TargetMachineID: "",
		},
	}

	if !reflect.DeepEqual(expectJobs, jobs) {
		t.Fatalf("Received unexpected collection of Jobs: %#v\nExpected: %#v", jobs, expectJobs)
	}

	a.unloadJob("foo.service")

	// This sucks, but we have to do it if Agent.unloadJob is going to spin
	// off the real work that matters in a goroutine
	time.Sleep(200)

	jobs, err = a.jobs()
	if err != nil {
		t.Fatalf("Failed calling Agent.jobs: %v", err)
	}

	expectJobs = map[string]*job.Job{}
	if !reflect.DeepEqual(expectJobs, jobs) {
		t.Fatalf("Received unexpected collection of Jobs: %#v\nExpected: %#v", jobs, expectJobs)
	}
}

func TestAgentLoadStartStopUnit(t *testing.T) {
	uManager := unit.NewFakeUnitManager()
	usGenerator := unit.NewUnitStateGenerator(uManager)
	fReg := registry.NewFakeRegistry()
	mach := &machine.FakeMachine{MachineState: machine.MachineState{ID: "XXX"}}
	a := New(uManager, usGenerator, fReg, mach, time.Second)

	u := newTestUnitFromUnitContents(t, "foo.service", "")

	err := a.loadUnit(u)
	if err != nil {
		t.Fatalf("Failed calling Agent.loadUnit: %v", err)
	}

	err = a.startUnit("foo.service")
	if err != nil {
		t.Fatalf("Failed starting unit foo.service: %v", err)
	}

	units, err := a.units()
	if err != nil {
		t.Fatalf("Failed calling Agent.units: %v", err)
	}

	jsLaunched := job.JobStateLaunched
	expectUnits := unitStates{
		"foo.service": unitState{
			state: jsLaunched,
		},
	}

	if !reflect.DeepEqual(expectUnits, units) {
		t.Fatalf("Received unexpected collection of Units: %#v\nExpected: %#v", units, expectUnits)
	}

	err = a.stopUnit("foo.service")
	if err != nil {
		t.Fatalf("Failed stopping unit foo.service: %v", err)
	}

	units, err = a.units()
	if err != nil {
		t.Fatalf("Failed calling Agent.units: %v", err)
	}

	jsLoaded := job.JobStateLoaded
	expectUnits = unitStates{
		"foo.service": unitState{
			state: jsLoaded,
		},
	}

	if !reflect.DeepEqual(expectUnits, units) {
		t.Fatalf("Received unexpected collection of Units: %#v\nExpected: %#v", units, expectUnits)
	}
}

func newFakeRegistryForCheckVersion(v string) registry.Registry {
	sv, err := semver.NewVersion(v)
	if err != nil {
		panic(err)
	}

	reg := registry.NewFakeRegistry()
	reg.SetLatestVersion(*sv)
	return reg
}

func newFakeRegistryForCommands(unitPrefix string, unitCount int, template bool) client.API {
	// clear machineStates for every invocation
	machineStates = nil
	machines := []machine.MachineState{
		newMachineState("c31e44e1-f858-436e-933e-59c642517860", "1.2.3.4", map[string]string{"ping": "pong"}),
		newMachineState("595989bb-cbb7-49ce-8726-722d6e157b4e", "5.6.7.8", map[string]string{"foo": "bar"}),
	}

	jobs := make([]job.Job, 0)
	appendJobsForTests(&jobs, machines[0], unitPrefix, unitCount, template)
	appendJobsForTests(&jobs, machines[1], unitPrefix, unitCount, template)

	states := make([]unit.UnitState, 0)
	if template {
		state := unit.UnitState{
			UnitName:    fmt.Sprintf("%[email protected]", unitPrefix),
			LoadState:   "loaded",
			ActiveState: "inactive",
			SubState:    "dead",
			MachineID:   machines[0].ID,
		}
		states = append(states, state)
		state.MachineID = machines[1].ID
		states = append(states, state)
	} else {
		for i := 1; i <= unitCount; i++ {
			state := unit.UnitState{
				UnitName:    fmt.Sprintf("%s%d.service", unitPrefix, i),
				LoadState:   "loaded",
				ActiveState: "active",
				SubState:    "listening",
				MachineID:   machines[0].ID,
			}
			states = append(states, state)
		}

		for i := 1; i <= unitCount; i++ {
			state := unit.UnitState{
				UnitName:    fmt.Sprintf("%s%d.service", unitPrefix, i),
				LoadState:   "loaded",
				ActiveState: "inactive",
				SubState:    "dead",
				MachineID:   machines[1].ID,
			}
			states = append(states, state)
		}
	}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(machines)
	reg.SetUnitStates(states)
	reg.SetJobs(jobs)

	return &client.RegistryClient{Registry: reg}
}

func newTestRegistryForListMachines() registry.Registry {
	m := []machine.MachineState{
		machine.MachineState{ID: "mnopqr"},
		machine.MachineState{ID: "abcdef"},
		machine.MachineState{ID: "ghijkl"},
	}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(m)

	return reg
}

func fakeMachinesSetup() (*machinesResource, *httptest.ResponseRecorder) {
	fr := registry.NewFakeRegistry()
	fr.SetMachines([]machine.MachineState{
		{ID: "XXX", PublicIP: "", Metadata: map[string]string{}},
		{ID: "YYY", PublicIP: "1.2.3.4", Metadata: map[string]string{"ping": "pong"}},
	})
	fAPI := &client.RegistryClient{Registry: fr}
	resource := &machinesResource{cAPI: fAPI, tokenLimit: testTokenLimit}
	rw := httptest.NewRecorder()

	return resource, rw
}

func newFakeRegistryForListUnits(jobs []job.Job) registry.Registry {
	j := []job.Job{*job.NewJob("pong.service", *unit.NewUnit("Echo"))}

	if jobs != nil {
		for _, job := range jobs {
			j = append(j, job)
		}
	}

	reg := registry.NewFakeRegistry()
	reg.SetJobs(j)

	return reg
}

func setupRegistryForStart(echoAttempts int) {
	m1 := machine.MachineState{
		ID:       "c31e44e1-f858-436e-933e-59c642517860",
		PublicIP: "1.2.3.4",
		Metadata: map[string]string{"ping": "pong"},
	}
	m2 := machine.MachineState{
		ID:       "595989bb-cbb7-49ce-8726-722d6e157b4e",
		PublicIP: "5.6.7.8",
		Metadata: map[string]string{"foo": "bar"},
	}
	m3 := machine.MachineState{
		ID:       "520983A8-FB9C-4A68-B49C-CED5BB2E9D08",
		Metadata: map[string]string{"foo": "bar"},
	}

	states := []unit.UnitState{
		unit.UnitState{
			UnitName:    "pong.service",
			LoadState:   "loaded",
			ActiveState: "active",
			SubState:    "listening",
			MachineID:   m1.ID,
		},
		unit.UnitState{
			UnitName:    "hello.service",
			LoadState:   "loaded",
			ActiveState: "inactive",
			SubState:    "dead",
			MachineID:   m2.ID,
		},
		unit.UnitState{
			UnitName:    "echo.service",
			LoadState:   "loaded",
			ActiveState: "inactive",
			SubState:    "dead",
			MachineID:   m2.ID,
		},
	}

	machines := []machine.MachineState{m1, m2, m3}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(machines)
	reg.SetUnitStates(states)

	cAPI = &client.RegistryClient{Registry: &BlockedFakeRegistry{EchoAttempts: echoAttempts, FakeRegistry: *reg}}
}

func TestUnitsListBadNextPageToken(t *testing.T) {
	fr := registry.NewFakeRegistry()
	resource := &unitsResource{fr, "/units"}
	rw := httptest.NewRecorder()
	req, err := http.NewRequest("GET", "http://example.com/units?nextPageToken=EwBMLg==", nil)
	if err != nil {
		t.Fatalf("Failed creating http.Request: %v", err)
	}

	resource.list(rw, req)

	err = assertErrorResponse(rw, http.StatusBadRequest)
	if err != nil {
		t.Error(err.Error())
	}
}

func newFakeRegistryForSsh() client.API {
	// clear machineStates for every invocation
	machineStates = nil
	machines := []machine.MachineState{
		newMachineState("c31e44e1-f858-436e-933e-59c642517860", "1.2.3.4", map[string]string{"ping": "pong"}),
		newMachineState("595989bb-cbb7-49ce-8726-722d6e157b4e", "5.6.7.8", map[string]string{"foo": "bar"}),
		newMachineState("hello.service", "8.7.6.5", map[string]string{"foo": "bar"}),
	}

	jobs := []job.Job{
		job.Job{Name: "j1.service", Unit: unit.UnitFile{}, TargetMachineID: machines[0].ID},
		job.Job{Name: "j2.service", Unit: unit.UnitFile{}, TargetMachineID: machines[1].ID},
		job.Job{Name: "hello.service", Unit: unit.UnitFile{}, TargetMachineID: machines[2].ID},
	}

	states := []unit.UnitState{
		unit.UnitState{
			UnitName:    "j1.service",
			LoadState:   "loaded",
			ActiveState: "active",
			SubState:    "listening",
			MachineID:   machines[0].ID,
		},
		unit.UnitState{
			UnitName:    "j2.service",
			LoadState:   "loaded",
			ActiveState: "inactive",
			SubState:    "dead",
			MachineID:   machines[1].ID,
		},
		unit.UnitState{
			UnitName:    "hello.service",
			LoadState:   "loaded",
			ActiveState: "inactive",
			SubState:    "dead",
			MachineID:   machines[2].ID,
		},
	}

	reg := registry.NewFakeRegistry()
	reg.SetMachines(machines)
	reg.SetUnitStates(states)
	reg.SetJobs(jobs)

	return &client.RegistryClient{Registry: reg}
}