mirror of
https://github.com/netbirdio/gvisor.git
synced 2026-05-22 17:12:49 -07:00
Move test server into a separate package
It will be used for tests in other packages. Updates #4805 PiperOrigin-RevId: 447076300
This commit is contained in:
Fabricio Voznika
committed by
gVisor bot
gVisor bot
parent
592fc1bb50
commit
71add37390
@@ -13,6 +13,7 @@ go_library(
|
||||
"//pkg/fd",
|
||||
"//pkg/log",
|
||||
"//pkg/sentry/seccheck",
|
||||
"//pkg/sentry/seccheck/checkers/remote/header",
|
||||
"//pkg/sentry/seccheck/points:points_go_proto",
|
||||
"@org_golang_google_protobuf//proto:go_default_library",
|
||||
"@org_golang_x_sys//unix:go_default_library",
|
||||
@@ -28,15 +29,13 @@ go_test(
|
||||
],
|
||||
library = ":remote",
|
||||
deps = [
|
||||
"//pkg/cleanup",
|
||||
"//pkg/fd",
|
||||
"//pkg/sentry/seccheck",
|
||||
"//pkg/sentry/seccheck/checkers/remote/test",
|
||||
"//pkg/sentry/seccheck/points:points_go_proto",
|
||||
"//pkg/sync",
|
||||
"//pkg/test/testutil",
|
||||
"@com_github_cenkalti_backoff//:go_default_library",
|
||||
"@org_golang_google_protobuf//proto:go_default_library",
|
||||
"@org_golang_google_protobuf//types/known/anypb:go_default_library",
|
||||
"@org_golang_x_sys//unix:go_default_library",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
load("//tools:defs.bzl", "go_library", "go_test")
|
||||
|
||||
package(licenses = ["notice"])
|
||||
|
||||
go_library(
|
||||
name = "header",
|
||||
srcs = ["header.go"],
|
||||
marshal = True,
|
||||
visibility = ["//:sandbox"],
|
||||
)
|
||||
|
||||
go_test(
|
||||
name = "header_test",
|
||||
size = "small",
|
||||
srcs = ["header_test.go"],
|
||||
library = ":header",
|
||||
)
|
||||
@@ -0,0 +1,44 @@
|
||||
// Copyright 2022 The gVisor Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Package header contains the message header used in the remote checker.
|
||||
package header
|
||||
|
||||
// HeaderStructSize size of header struct in bytes.
|
||||
const HeaderStructSize = 8
|
||||
|
||||
// Header is used to describe the message being sent to the remote process.
|
||||
//
|
||||
// 0 --------- 16 ---------- 32 ----------- 64 -----------+
|
||||
// | HeaderSize | MessageType | DroppedCount | Payload... |
|
||||
// +---- 16 ----+---- 16 -----+----- 32 -----+------------+
|
||||
//
|
||||
// +marshal
|
||||
type Header struct {
|
||||
// HeaderSize is the size of the header in bytes. The payload comes
|
||||
// immediatelly after the header. The length is needed to allow the header to
|
||||
// expand in the future without breaking remotes that do not yet understand
|
||||
// the new fields.
|
||||
HeaderSize uint16
|
||||
|
||||
// MessageType describes the payload. It must be one of the pb.MessageType
|
||||
// values and determine how the payload is interpreted. This is more efficient
|
||||
// than using protobuf.Any because Any uses the full protobuf name to identify
|
||||
// the type.
|
||||
MessageType uint16
|
||||
|
||||
// DroppedCount is the number of points that failed to be written and had to
|
||||
// be dropped. It wraps around after max(uint32).
|
||||
DroppedCount uint32
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
// Copyright 2022 The gVisor Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package header
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestHeaderSize(t *testing.T) {
|
||||
hdr := Header{}
|
||||
if want, got := hdr.SizeBytes(), HeaderStructSize; want != got {
|
||||
t.Errorf("wrong const header size, want: %v, got: %v", want, got)
|
||||
}
|
||||
}
|
||||
@@ -20,14 +20,14 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"gvisor.dev/gvisor/pkg/log"
|
||||
|
||||
"golang.org/x/sys/unix"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"gvisor.dev/gvisor/pkg/cleanup"
|
||||
"gvisor.dev/gvisor/pkg/context"
|
||||
"gvisor.dev/gvisor/pkg/fd"
|
||||
"gvisor.dev/gvisor/pkg/log"
|
||||
"gvisor.dev/gvisor/pkg/sentry/seccheck"
|
||||
"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/header"
|
||||
|
||||
pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
|
||||
)
|
||||
@@ -101,45 +101,17 @@ func New(_ map[string]interface{}, endpoint *fd.FD) (seccheck.Checker, error) {
|
||||
return &Remote{endpoint: endpoint}, nil
|
||||
}
|
||||
|
||||
// Header is used to describe the message being sent to the remote process.
|
||||
//
|
||||
// 0 --------- 16 ---------- 32 ----------- 64 -----------+
|
||||
// | HeaderSize | MessageType | DroppedCount | Payload... |
|
||||
// +---- 16 ----+---- 16 -----+----- 32 -----+------------+
|
||||
//
|
||||
// +marshal
|
||||
type Header struct {
|
||||
// HeaderSize is the size of the header in bytes. The payload comes
|
||||
// immediatelly after the header. The length is needed to allow the header to
|
||||
// expand in the future without breaking remotes that do not yet understand
|
||||
// the new fields.
|
||||
HeaderSize uint16
|
||||
|
||||
// MessageType describes the payload. It must be one of the pb.MessageType
|
||||
// values and determine how the payload is interpreted. This is more efficient
|
||||
// than using protobuf.Any because Any uses the full protobuf name to identify
|
||||
// the type.
|
||||
MessageType uint16
|
||||
|
||||
// DroppedCount is the number of points that failed to be written and had to
|
||||
// be dropped. It wraps around after max(uint32).
|
||||
DroppedCount uint32
|
||||
}
|
||||
|
||||
// headerStructSize size of header struct in bytes.
|
||||
const headerStructSize = 8
|
||||
|
||||
func (r *Remote) write(msg proto.Message, msgType pb.MessageType) {
|
||||
out, err := proto.Marshal(msg)
|
||||
if err != nil {
|
||||
log.Debugf("Marshal(%+v): %v", msg, err)
|
||||
return
|
||||
}
|
||||
hdr := Header{
|
||||
HeaderSize: uint16(headerStructSize),
|
||||
hdr := header.Header{
|
||||
HeaderSize: uint16(header.HeaderStructSize),
|
||||
MessageType: uint16(msgType),
|
||||
}
|
||||
var hdrOut [headerStructSize]byte
|
||||
var hdrOut [header.HeaderStructSize]byte
|
||||
hdr.MarshalUnsafe(hdrOut[:])
|
||||
|
||||
// TODO(gvisor.dev/issue/4805): Change to non-blocking write. Count as dropped
|
||||
|
||||
@@ -17,7 +17,6 @@ package remote
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
@@ -26,14 +25,12 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/cenkalti/backoff"
|
||||
"golang.org/x/sys/unix"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/anypb"
|
||||
"gvisor.dev/gvisor/pkg/cleanup"
|
||||
"gvisor.dev/gvisor/pkg/fd"
|
||||
"gvisor.dev/gvisor/pkg/sentry/seccheck"
|
||||
"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/test"
|
||||
pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
|
||||
"gvisor.dev/gvisor/pkg/sync"
|
||||
"gvisor.dev/gvisor/pkg/test/testutil"
|
||||
)
|
||||
|
||||
@@ -91,133 +88,14 @@ func (s *exampleServer) stop() {
|
||||
_ = os.Remove(s.path)
|
||||
}
|
||||
|
||||
type server struct {
|
||||
path string
|
||||
fd *fd.FD
|
||||
stopCh chan struct{}
|
||||
|
||||
mu sync.Mutex
|
||||
// +checklocks:mu
|
||||
points []message
|
||||
}
|
||||
|
||||
type message struct {
|
||||
msgType pb.MessageType
|
||||
msg []byte
|
||||
}
|
||||
|
||||
func newServer() (*server, error) {
|
||||
dir, err := ioutil.TempDir(os.TempDir(), "remote")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
server, err := newServerPath(filepath.Join(dir, "remote.sock"))
|
||||
if err != nil {
|
||||
_ = os.RemoveAll(dir)
|
||||
return nil, err
|
||||
}
|
||||
return server, nil
|
||||
}
|
||||
|
||||
func newServerPath(path string) (*server, error) {
|
||||
socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
|
||||
}
|
||||
cu := cleanup.Make(func() {
|
||||
_ = unix.Close(socket)
|
||||
})
|
||||
defer cu.Clean()
|
||||
|
||||
sa := &unix.SockaddrUnix{Name: path}
|
||||
if err := unix.Bind(socket, sa); err != nil {
|
||||
return nil, fmt.Errorf("bind(%q): %w", path, err)
|
||||
}
|
||||
if err := unix.Listen(socket, 5); err != nil {
|
||||
return nil, fmt.Errorf("listen(): %w", err)
|
||||
}
|
||||
|
||||
server := &server{
|
||||
path: path,
|
||||
fd: fd.New(socket),
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
go server.run()
|
||||
cu.Release()
|
||||
return server, nil
|
||||
}
|
||||
|
||||
func (s *server) run() {
|
||||
defer func() {
|
||||
s.stopCh <- struct{}{}
|
||||
}()
|
||||
for {
|
||||
client, _, err := unix.Accept(s.fd.FD())
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
go s.handleClient(client)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *server) handleClient(client int) {
|
||||
defer unix.Close(client)
|
||||
|
||||
var buf = make([]byte, 1024*1024)
|
||||
for {
|
||||
read, err := unix.Read(client, buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if read == 0 {
|
||||
return
|
||||
}
|
||||
if read <= headerStructSize {
|
||||
panic("invalid message")
|
||||
}
|
||||
hdr := Header{}
|
||||
hdr.UnmarshalUnsafe(buf[0:headerStructSize])
|
||||
msg := message{
|
||||
msgType: pb.MessageType(hdr.MessageType),
|
||||
msg: buf[hdr.HeaderSize:read],
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.points = append(s.points, msg)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *server) count() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return len(s.points)
|
||||
}
|
||||
|
||||
func (s *server) getPoints() []message {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cpy := make([]message, len(s.points))
|
||||
copy(cpy, s.points)
|
||||
return cpy
|
||||
}
|
||||
|
||||
func (s *server) wait() {
|
||||
<-s.stopCh
|
||||
}
|
||||
|
||||
func (s *server) close() {
|
||||
_ = s.fd.Close()
|
||||
_ = os.Remove(s.path)
|
||||
}
|
||||
|
||||
func TestBasic(t *testing.T) {
|
||||
server, err := newServer()
|
||||
server, err := test.NewServer()
|
||||
if err != nil {
|
||||
t.Fatalf("newServer(): %v", err)
|
||||
}
|
||||
defer server.close()
|
||||
defer server.Close()
|
||||
|
||||
endpoint, err := setup(server.path)
|
||||
endpoint, err := setup(server.Path)
|
||||
if err != nil {
|
||||
t.Fatalf("setup(): %v", err)
|
||||
}
|
||||
@@ -238,27 +116,22 @@ func TestBasic(t *testing.T) {
|
||||
t.Fatalf("ExitNotifyParent: %v", err)
|
||||
}
|
||||
|
||||
testutil.Poll(func() error {
|
||||
if server.count() == 0 {
|
||||
return fmt.Errorf("waiting for points to arrive")
|
||||
}
|
||||
return nil
|
||||
}, 5*time.Second)
|
||||
if want, got := 1, server.count(); want != got {
|
||||
t.Errorf("wrong number of points, want: %d, got: %d", want, got)
|
||||
}
|
||||
pt := server.getPoints()[0]
|
||||
|
||||
if want := pb.MessageType_MESSAGE_SENTRY_EXIT_NOTIFY_PARENT; pt.msgType != want {
|
||||
t.Errorf("wrong message type, want: %v, got: %v", want, pt.msgType)
|
||||
server.WaitForCount(1)
|
||||
pt := server.GetPoints()[0]
|
||||
if want := pb.MessageType_MESSAGE_SENTRY_EXIT_NOTIFY_PARENT; pt.MsgType != want {
|
||||
t.Errorf("wrong message type, want: %v, got: %v", want, pt.MsgType)
|
||||
}
|
||||
got := &pb.ExitNotifyParentInfo{}
|
||||
if err := proto.Unmarshal(pt.msg, got); err != nil {
|
||||
if err := proto.Unmarshal(pt.Msg, got); err != nil {
|
||||
t.Errorf("proto.Unmarshal(ExitNotifyParentInfo): %v", err)
|
||||
}
|
||||
if !proto.Equal(info, got) {
|
||||
t.Errorf("Received point is different, want: %+v, got: %+v", info, got)
|
||||
}
|
||||
// Check that no more points were received.
|
||||
if want, got := 1, server.Count(); want != got {
|
||||
t.Errorf("wrong number of points, want: %d, got: %d", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
// Test that the example C++ server works. It's easier to test from here and
|
||||
@@ -304,13 +177,6 @@ func TestExample(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeaderSize(t *testing.T) {
|
||||
hdr := Header{}
|
||||
if want, got := hdr.SizeBytes(), hdr.SizeBytes(); want != got {
|
||||
t.Errorf("wrong const header size, want: %v, got: %v", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSmall(t *testing.B) {
|
||||
// Run server in a separate process just to isolate it as much as possible.
|
||||
server, err := newExampleServer(false)
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
load("//tools:defs.bzl", "go_library")
|
||||
|
||||
package(licenses = ["notice"])
|
||||
|
||||
go_library(
|
||||
name = "test",
|
||||
testonly = True,
|
||||
srcs = ["server.go"],
|
||||
visibility = ["//:sandbox"],
|
||||
deps = [
|
||||
"//pkg/cleanup",
|
||||
"//pkg/log",
|
||||
"//pkg/sentry/seccheck/checkers/remote/header",
|
||||
"//pkg/sentry/seccheck/points:points_go_proto",
|
||||
"//pkg/sync",
|
||||
"//pkg/test/testutil",
|
||||
"//pkg/unet",
|
||||
"@org_golang_x_sys//unix:go_default_library",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,201 @@
|
||||
// Copyright 2022 The gVisor Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Package test provides functionality used to test the remote checker.
|
||||
package test
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"golang.org/x/sys/unix"
|
||||
"gvisor.dev/gvisor/pkg/cleanup"
|
||||
"gvisor.dev/gvisor/pkg/log"
|
||||
"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/header"
|
||||
pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
|
||||
"gvisor.dev/gvisor/pkg/sync"
|
||||
"gvisor.dev/gvisor/pkg/test/testutil"
|
||||
"gvisor.dev/gvisor/pkg/unet"
|
||||
)
|
||||
|
||||
// Server is the counterpart to the checkers.Remote. It receives connections
|
||||
// remote checkers and stores all points that it receives.
|
||||
type Server struct {
|
||||
Path string
|
||||
socket *unet.ServerSocket
|
||||
|
||||
mu sync.Mutex
|
||||
|
||||
// +checklocks:mu
|
||||
clients []*unet.Socket
|
||||
|
||||
// +checklocks:mu
|
||||
points []Message
|
||||
}
|
||||
|
||||
// Message corresponds to a single message sent from checkers.Remote.
|
||||
type Message struct {
|
||||
// MsgType indicates what is the type of Msg.
|
||||
MsgType pb.MessageType
|
||||
// Msg is the payload to the message that can be decoded using MsgType.
|
||||
Msg []byte
|
||||
}
|
||||
|
||||
// NewServer creates a new server that listens to a UDS that it creates under
|
||||
// os.TempDir.
|
||||
func NewServer() (*Server, error) {
|
||||
dir, err := ioutil.TempDir(os.TempDir(), "remote")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
server, err := newServerPath(filepath.Join(dir, "remote.sock"))
|
||||
if err != nil {
|
||||
_ = os.RemoveAll(dir)
|
||||
return nil, err
|
||||
}
|
||||
return server, nil
|
||||
}
|
||||
|
||||
func newServerPath(path string) (*Server, error) {
|
||||
socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
|
||||
}
|
||||
cu := cleanup.Make(func() {
|
||||
_ = unix.Close(socket)
|
||||
})
|
||||
defer cu.Clean()
|
||||
|
||||
sa := &unix.SockaddrUnix{Name: path}
|
||||
if err := unix.Bind(socket, sa); err != nil {
|
||||
return nil, fmt.Errorf("bind(%q): %w", path, err)
|
||||
}
|
||||
|
||||
ss, err := unet.NewServerSocket(socket)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cu.Add(func() { ss.Close() })
|
||||
|
||||
if err := ss.Listen(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
server := &Server{
|
||||
Path: path,
|
||||
socket: ss,
|
||||
}
|
||||
go server.run()
|
||||
cu.Release()
|
||||
return server, nil
|
||||
}
|
||||
|
||||
func (s *Server) run() {
|
||||
for {
|
||||
client, err := s.socket.Accept()
|
||||
if err != nil {
|
||||
// EBADF returns when the socket closes.
|
||||
if !errors.Is(err, unix.EBADF) {
|
||||
log.Warningf("socket.Accept(): %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.clients = append(s.clients, client)
|
||||
s.mu.Unlock()
|
||||
go s.handleClient(client)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleClient(client *unet.Socket) {
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
for i, c := range s.clients {
|
||||
if c == client {
|
||||
s.clients = append(s.clients[:i], s.clients[i+1:]...)
|
||||
break
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
_ = client.Close()
|
||||
}()
|
||||
|
||||
var buf = make([]byte, 1024*1024)
|
||||
for {
|
||||
read, err := client.Read(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if read == 0 {
|
||||
return
|
||||
}
|
||||
if read < header.HeaderStructSize {
|
||||
panic("invalid message")
|
||||
}
|
||||
hdr := header.Header{}
|
||||
hdr.UnmarshalUnsafe(buf[0:header.HeaderStructSize])
|
||||
if read < int(hdr.HeaderSize) {
|
||||
panic(fmt.Sprintf("message truncated, header size: %d, readL %d", hdr.HeaderSize, read))
|
||||
}
|
||||
msg := Message{
|
||||
MsgType: pb.MessageType(hdr.MessageType),
|
||||
Msg: buf[hdr.HeaderSize:read],
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.points = append(s.points, msg)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// Count return the number of points it has received.
|
||||
func (s *Server) Count() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return len(s.points)
|
||||
}
|
||||
|
||||
// GetPoints returns all points that it has received.
|
||||
func (s *Server) GetPoints() []Message {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cpy := make([]Message, len(s.points))
|
||||
copy(cpy, s.points)
|
||||
return cpy
|
||||
}
|
||||
|
||||
// Close stops listenning and closes all connections.
|
||||
func (s *Server) Close() {
|
||||
_ = s.socket.Close()
|
||||
s.mu.Lock()
|
||||
for _, client := range s.clients {
|
||||
_ = client.Close()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
_ = os.Remove(s.Path)
|
||||
}
|
||||
|
||||
// WaitForCount waits for the number of points to reach the desired number for
|
||||
// 5 seconds. It fails if not received in time.
|
||||
func (s *Server) WaitForCount(count int) error {
|
||||
return testutil.Poll(func() error {
|
||||
if got := s.Count(); got < count {
|
||||
return fmt.Errorf("waiting for points %d to arrive, received %d", count, got)
|
||||
}
|
||||
return nil
|
||||
}, 5*time.Second)
|
||||
}
|
||||
Reference in New Issue
Block a user