Ejemplo n.º 1
0
func testAgent(t *testing.T) *agentWrapper {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	lw := agent.NewLogWriter(512)
	mult := io.MultiWriter(os.Stderr, lw)

	conf := nextConfig()

	dir, err := ioutil.TempDir("", "agent")
	if err != nil {
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}
	conf.DataDir = dir

	a, err := agent.Create(conf, lw)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	rpc := agent.NewAgentRPC(a, l, mult, lw)
	return &agentWrapper{
		dir:    dir,
		config: conf,
		agent:  a,
		rpc:    rpc,
		addr:   l.Addr().String(),
	}
}
Ejemplo n.º 2
0
func testAgentWithConfig(t *testing.T, cb func(c *agent.Config)) *agentWrapper {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	lw := logger.NewLogWriter(512)
	mult := io.MultiWriter(os.Stderr, lw)

	conf := nextConfig()
	cb(conf)

	dir, err := ioutil.TempDir("", "agent")
	if err != nil {
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}
	conf.DataDir = dir

	a, err := agent.Create(conf, lw, nil, nil)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	rpc := agent.NewAgentRPC(a, l, mult, lw)

	conf.Addresses.HTTP = "127.0.0.1"
	httpAddr := fmt.Sprintf("127.0.0.1:%d", conf.Ports.HTTP)
	http, err := agent.NewHTTPServers(a, conf, os.Stderr)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	if http == nil || len(http) == 0 {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("Could not create HTTP server to listen on: %s", httpAddr))
	}

	return &agentWrapper{
		dir:      dir,
		config:   conf,
		agent:    a,
		rpc:      rpc,
		http:     http[0],
		addr:     l.Addr().String(),
		httpAddr: httpAddr,
	}
}
func (s *Server) ServeTCP() {
	mockAgent := new(FakeAgentBackend)
	if s.FailStatsEndpoint {
		mockAgent.StatsReturns(map[string]map[string]string{
			"raft": {
				"commit_index":   "5",
				"last_log_index": "2",
			},
		})
	}

	agentRPCServer := agent.NewAgentRPC(mockAgent, s.TCPListener, os.Stderr, agent.NewLogWriter(42))

	var (
		useKeyCallCount     int
		installKeyCallCount int
		leaveCallCount      int
		statsCallCount      int
	)

	for {
		switch {
		case mockAgent.UseKeyCallCount() > useKeyCallCount:
			useKeyCallCount++
			s.OutputWriter.UseKeyCalled()
		case mockAgent.InstallKeyCallCount() > installKeyCallCount:
			installKeyCallCount++
			s.OutputWriter.InstallKeyCalled()
		case mockAgent.LeaveCallCount() > leaveCallCount:
			leaveCallCount++
			s.OutputWriter.LeaveCalled()
			agentRPCServer.Shutdown()
			s.DidLeave = true
		case mockAgent.StatsCallCount() > statsCallCount:
			statsCallCount++
			s.OutputWriter.StatsCalled()
		}

		time.Sleep(10 * time.Millisecond)
	}
}
Ejemplo n.º 4
0
func (sl ServerListener) Serve() {
	httpListener, err := net.Listen("tcp", sl.HTTPAddr)
	if err != nil {
		panic(err)
	}

	triggerClose := make(chan struct{})

	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		var members []api.AgentMember
		for _, member := range sl.Members {
			members = append(members, api.AgentMember{
				Addr: member,
			})
		}
		json.NewEncoder(w).Encode(members)
		if !sl.StayAlive {
			triggerClose <- struct{}{}
		}
	})

	server := &http.Server{
		Addr:    sl.HTTPAddr,
		Handler: mux,
	}

	go server.Serve(httpListener)

	var mockAgent *FakeAgentBackend
	var agentRPCServer *agent.AgentRPC
	if sl.TCPAddr != "" {

		tcpAddr, err := net.ResolveTCPAddr("tcp", sl.TCPAddr)
		if err != nil {
			panic(err)
		}

		tcpListener, err := net.ListenTCP("tcp", tcpAddr)
		if err != nil {
			panic(err)
		}

		mockAgent = new(FakeAgentBackend)
		agentRPCServer = agent.NewAgentRPC(mockAgent, tcpListener,
			os.Stderr, agent.NewLogWriter(42))

		go func() {
			var useKeyCalled bool
			for {
				if mockAgent.UseKeyCallCount() > 0 && !useKeyCalled {
					fmt.Println("UseKey called")
					useKeyCalled = true
				}

				time.Sleep(1 * time.Second)
			}
		}()

		go func() {
			var leaveCalled bool
			for {
				if mockAgent.LeaveCallCount() > 0 && !leaveCalled {
					fmt.Println("Leave called")
					leaveCalled = true
					triggerClose <- struct{}{}
					triggerClose <- struct{}{}
				}

				time.Sleep(1 * time.Second)
			}
		}()
	}

	<-triggerClose
	<-triggerClose
	time.Sleep(1 * time.Second)
	httpListener.Close()
	if agentRPCServer != nil {
		agentRPCServer.Shutdown()
	}
}