go/src/vanadium/discovery/internal/discovery_test.go

package internal

import (
	"reflect"
	"testing"

	"v.io/v23/context"
	"v.io/v23/discovery"
	"v.io/v23/security/access"
	idiscovery "v.io/x/ref/lib/discovery"
	vtest "v.io/x/ref/test"

	_ "v.io/x/ref/runtime/factories/generic"

	mojom "mojom/vanadium/discovery"
)

type mockAdv struct {
	s discovery.Service
}

type discoveryMock struct {
	trigger  *idiscovery.Trigger
	id       int64
	services map[int64]discovery.Service
	// An item will be put in deleteCh when something has been deleted.
	deleteCh chan struct{}
}

func (d *discoveryMock) Advertise(ctx *context.T, s discovery.Service, perms access.Permissions) error {
	currId := d.id
	d.services[currId] = s
	d.id++
	c := func() {
		delete(d.services, currId)
		d.deleteCh <- struct{}{}
	}
	d.trigger.Add(c, ctx.Done())
	return nil
}

func (*discoveryMock) Scan(ctx *context.T, query string) (<-chan discovery.Update, error) {
	return nil, nil
}

func compare(t *testing.T, want discovery.Service, got mojom.Service) {
	mwant := v2mService(want)
	if !reflect.DeepEqual(mwant, got) {
		t.Errorf("Got %#v want %#v", got, want)
	}
}

func TestAdvertising(t *testing.T) {
	ctx, shutdown := vtest.V23Init()
	defer shutdown()
	mock := &discoveryMock{
		trigger:  idiscovery.NewTrigger(),
		services: map[int64]discovery.Service{},
		deleteCh: make(chan struct{}),
	}
	s := NewDiscoveryService(ctx, mock)
	testService := mojom.Service{
		InterfaceName: "v.io/v23/discovery.T",
		Attrs: map[string]string{
			"key1": "value1",
			"key2": "value2",
		},
		Addrs: []string{"addr1", "addr2"},
	}
	id, e1, e2 := s.Advertise(testService, nil)

	if e1 != nil || e2 != nil {
		t.Fatalf("Failed to start service: %v, %v", e1, e2)
	}
	if len(mock.services) != 1 {
		t.Errorf("service missing in mock")
	}

	for _, service := range mock.services {
		compare(t, service, testService)
	}

	testService2 := mojom.Service{
		InterfaceName: "v.io/v23/naming.T",
		Attrs: map[string]string{
			"key1": "value1",
			"key2": "value2",
		},
		Addrs: []string{"addr1", "addr2"},
	}

	_, e1, e2 = s.Advertise(testService2, nil)
	if e1 != nil || e2 != nil {
		t.Fatalf("Failed to start service: %v, %v", e1, e2)
	}

	s.StopAdvertising(id)
	// Wait for the deletion to finish.
	<-mock.deleteCh
	if len(mock.services) != 1 {
		t.Errorf("service should have been removed")
	}

	for _, service := range mock.services {
		compare(t, service, testService2)
	}

	s.Stop()
	<-mock.deleteCh
	if len(mock.services) != 0 {
		t.Errorf("service should have been removed")
	}
}