traefik/pkg/pilot/pilot_test.go

125 lines
2.9 KiB
Go
Raw Normal View History

package pilot
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/traefik/traefik/v2/pkg/config/runtime"
"github.com/traefik/traefik/v2/pkg/metrics"
"github.com/traefik/traefik/v2/pkg/safe"
)
func TestTick(t *testing.T) {
receivedConfig := make(chan bool)
mux := http.NewServeMux()
server := httptest.NewServer(mux)
t.Cleanup(server.Close)
mux.HandleFunc("/", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
err := json.NewEncoder(rw).Encode(instanceInfo{ID: "123"})
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
})
mux.HandleFunc("/command", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
receivedConfig <- true
})
pilot := New("token", metrics.RegisterPilot(), safe.NewPool(context.Background()))
pilot.client.baseURL = server.URL
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
go pilot.Tick(ctx)
pilot.SetRuntimeConfiguration(&runtime.Configuration{})
pilot.SetRuntimeConfiguration(&runtime.Configuration{})
select {
case <-time.Tick(10 * time.Second):
t.Fatal("Timeout")
case <-receivedConfig:
return
}
}
func TestClient_SendConfiguration(t *testing.T) {
myToken := "myToken"
mux := http.NewServeMux()
server := httptest.NewServer(mux)
t.Cleanup(server.Close)
mux.HandleFunc("/", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
tk := req.Header.Get(tokenHeader)
if tk != myToken {
http.Error(rw, fmt.Sprintf("invalid token: %s", tk), http.StatusUnauthorized)
}
err := json.NewEncoder(rw).Encode(instanceInfo{ID: "123"})
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
})
mux.HandleFunc("/command", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
tk := req.Header.Get(tokenHeader)
if tk != myToken {
http.Error(rw, fmt.Sprintf("invalid token: %s", tk), http.StatusUnauthorized)
}
defer req.Body.Close()
info := &instanceInfo{}
err := json.NewDecoder(req.Body).Decode(info)
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
if info.ID != "123" {
http.Error(rw, fmt.Sprintf("invalid ID: %s", info.ID), http.StatusBadRequest)
}
})
client := client{
baseURL: server.URL,
httpClient: http.DefaultClient,
token: myToken,
}
err := client.SendData(context.Background(), RunTimeRepresentation{}, []metrics.PilotMetric{})
require.NoError(t, err)
}