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(),
}
}
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 := agent.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)
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, nil, 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)
}
}
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()
}
}