blob: 738e495d8403fe7882d835416d127d408339d700 [file] [log] [blame]
// Copyright 2016 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package testutil defines helpers for discovery tests.
package testutil
import (
"bytes"
"fmt"
"reflect"
"runtime"
"sort"
"time"
"v.io/v23"
"v.io/v23/context"
"v.io/v23/discovery"
"v.io/v23/security"
"v.io/x/ref/lib/security/bcrypter"
"v.io/x/ref/test/testutil"
)
// WithPrincipal creates a new principal with an extension of the default
// blessing of the principal in the returned context.
func WithPrincipal(ctx *context.T, extension string) (*context.T, error) {
idp := testutil.IDProviderFromPrincipal(v23.GetPrincipal(ctx))
ctx, err := v23.WithPrincipal(ctx, testutil.NewPrincipal())
if err != nil {
return nil, err
}
if err := idp.Bless(v23.GetPrincipal(ctx), extension); err != nil {
return nil, err
}
return ctx, nil
}
// WithPrivateKey creates a new principal with the provided blessing and
// stores the private key of the blessing in the returned context.
func WithPrivateKey(ctx *context.T, root *bcrypter.Root, blessing string) (*context.T, error) {
ctx, err := v23.WithPrincipal(ctx, testutil.NewPrincipal(blessing))
if err != nil {
return nil, err
}
key, err := root.Extract(ctx, blessing)
if err != nil {
return nil, err
}
crypter := bcrypter.NewCrypter()
ctx = bcrypter.WithCrypter(ctx, crypter)
if err := crypter.AddKey(ctx, key); err != nil {
return nil, err
}
return ctx, nil
}
func Advertise(ctx *context.T, d discovery.T, visibility []security.BlessingPattern, ads ...*discovery.Advertisement) (func(), error) {
ctx, cancel := context.WithCancel(ctx)
dones := make([]<-chan struct{}, 0, len(ads))
for _, ad := range ads {
done, err := d.Advertise(ctx, ad, visibility)
if err != nil {
cancel()
return nil, fmt.Errorf("Advertise failed: %v", err)
}
dones = append(dones, done)
}
stop := func() {
cancel()
for _, done := range dones {
<-done
}
}
return stop, nil
}
func Scan(ctx *context.T, d discovery.T, query string) (<-chan discovery.Update, func(), error) {
ctx, cancel := context.WithCancel(ctx)
scanCh, err := d.Scan(ctx, query)
if err != nil {
return nil, nil, fmt.Errorf("Scan failed: %v", err)
}
stop := func() {
cancel()
for range scanCh {
}
}
return scanCh, stop, err
}
func ScanAndMatch(ctx *context.T, d discovery.T, query string, wants ...discovery.Advertisement) error {
const timeout = 3 * time.Minute
var updates []discovery.Update
for now := time.Now(); time.Since(now) < timeout; {
var err error
updates, err = doScan(ctx, d, query, len(wants))
if err != nil {
return err
}
if MatchFound(ctx, updates, wants...) {
return nil
}
runtime.Gosched()
}
return fmt.Errorf("Match failed; got %v, but wanted %v", updates, wants)
}
func doScan(ctx *context.T, d discovery.T, query string, expectedUpdates int) ([]discovery.Update, error) {
scanCh, stop, err := Scan(ctx, d, query)
if err != nil {
return nil, err
}
defer stop()
updates := make([]discovery.Update, 0, expectedUpdates)
for {
var timer <-chan time.Time
if len(updates) >= expectedUpdates {
timer = time.After(5 * time.Millisecond)
}
select {
case update := <-scanCh:
updates = append(updates, update)
case <-timer:
return updates, nil
}
}
}
func MatchFound(ctx *context.T, updates []discovery.Update, wants ...discovery.Advertisement) bool {
return match(ctx, updates, false, wants...)
}
func MatchLost(ctx *context.T, updates []discovery.Update, wants ...discovery.Advertisement) bool {
return match(ctx, updates, true, wants...)
}
func match(ctx *context.T, updates []discovery.Update, lost bool, wants ...discovery.Advertisement) bool {
updateMap := make(map[discovery.AdId]discovery.Update)
for _, update := range updates {
updateMap[update.Id()] = update
}
for _, want := range wants {
update := updateMap[want.Id]
if update == nil {
return false
}
if update.IsLost() != lost {
return false
}
if !UpdateEqual(ctx, update, want) {
return false
}
delete(updateMap, update.Id())
}
return len(updateMap) == 0
}
func UpdateEqual(ctx *context.T, update discovery.Update, ad discovery.Advertisement) bool {
if update.Id() != ad.Id {
return false
}
if update.InterfaceName() != ad.InterfaceName {
return false
}
// Not that addresses may be in different orders.
a1 := append([]string(nil), update.Addresses()...)
a2 := append([]string(nil), ad.Addresses...)
sort.Strings(a1)
sort.Strings(a2)
if !reflect.DeepEqual(a1, a2) {
return false
}
for k, v := range ad.Attributes {
if update.Attribute(k) != v {
return false
}
}
for k, v := range ad.Attachments {
r, ok := <-update.Attachment(ctx, k)
if !ok {
return false
}
if r.Error != nil {
return false
}
if !bytes.Equal(r.Data, v) {
return false
}
}
return true
}