When possible downstream workers will include their worker id information in the connections they establish with the upstream using the nodeenrollment library. The tracking listener now tracks connections that are connected in this same way instead of just connections of type *tls.Conn.
66 lines
1.7 KiB
Go
66 lines
1.7 KiB
Go
// Copyright (c) HashiCorp, Inc.
|
|
// SPDX-License-Identifier: MPL-2.0
|
|
|
|
package cluster
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/protobuf/types/known/structpb"
|
|
)
|
|
|
|
type testConn struct {
|
|
net.Conn
|
|
st *structpb.Struct
|
|
}
|
|
|
|
func (c *testConn) State() *structpb.Struct {
|
|
return c.st
|
|
}
|
|
|
|
func TestWorkerConnectionInfo(t *testing.T) {
|
|
connInfo := []*WorkerConnectionInfo{
|
|
{WorkerId: "something"},
|
|
{UpstreamAddress: "something"},
|
|
{BoundaryVersion: "something"},
|
|
{WorkerId: "something", UpstreamAddress: "something else", BoundaryVersion: "something yet more"},
|
|
}
|
|
for _, c := range connInfo {
|
|
st, err := c.AsConnectionStateStruct()
|
|
assert.NoError(t, err)
|
|
|
|
got, err := GetWorkerInfoFromStateMap(context.Background(), &testConn{st: st})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, c, got)
|
|
}
|
|
}
|
|
|
|
func TestWorkerConnectionInfo_errors(t *testing.T) {
|
|
t.Run("wrong connection type", func(t *testing.T) {
|
|
c, _ := net.Pipe()
|
|
got, err := GetWorkerInfoFromStateMap(context.Background(), c)
|
|
assert.ErrorContains(t, err, "connection does not provide a State structpb.Struct")
|
|
assert.Nil(t, got)
|
|
})
|
|
t.Run("no state on connection", func(t *testing.T) {
|
|
got, err := GetWorkerInfoFromStateMap(context.Background(), &testConn{})
|
|
assert.ErrorContains(t, err, "no state attached")
|
|
assert.Nil(t, got)
|
|
})
|
|
|
|
t.Run("wrong type in struct", func(t *testing.T) {
|
|
st, err := structpb.NewStruct(map[string]interface{}{
|
|
"worker_id": 1.2480,
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, st)
|
|
got, err := GetWorkerInfoFromStateMap(context.Background(), &testConn{st: st})
|
|
assert.ErrorContains(t, err, "'worker_id' expected type 'string'")
|
|
assert.Nil(t, got)
|
|
})
|
|
}
|