From fbe41a04fd6a4c4367aed51245bd3ad0b9513a04 Mon Sep 17 00:00:00 2001
From: Todd <github@quaddi.com>
Date: Thu, 5 Oct 2023 10:57:54 -0700
Subject: [PATCH] fix search test with boundary token check (#3799)

---
 internal/cmd/commands/daemon/testing.go     | 11 ++---------
 internal/cmd/commands/search/search_test.go | 15 ++++++++-------
 2 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/internal/cmd/commands/daemon/testing.go b/internal/cmd/commands/daemon/testing.go
index 1f0db437fa..e22b123206 100644
--- a/internal/cmd/commands/daemon/testing.go
+++ b/internal/cmd/commands/daemon/testing.go
@@ -5,7 +5,6 @@ package daemon
 
 import (
 	"context"
-	stdErrors "errors"
 	"io"
 	"sync"
 	"testing"
@@ -71,10 +70,10 @@ func (s *TestServer) Serve(t *testing.T, opt ...Option) error {
 
 // AddResources adds targets to the cache for the provided address, token name,
 // and keyring type. They token info must already be known to the server.
-func (s *TestServer) AddResources(t *testing.T, p *authtokens.AuthToken, tars []*targets.Target, sess []*sessions.Session) {
+func (s *TestServer) AddResources(t *testing.T, p *authtokens.AuthToken, tars []*targets.Target, sess []*sessions.Session, atReadFn cache.BoundaryTokenReaderFn) {
 	t.Helper()
 	ctx := context.Background()
-	r, err := cache.NewRepository(ctx, s.cacheServer.store, &sync.Map{}, s.cmd.ReadTokenFromKeyring, unimplementedAuthTokenReader)
+	r, err := cache.NewRepository(ctx, s.cacheServer.store, &sync.Map{}, s.cmd.ReadTokenFromKeyring, atReadFn)
 	require.NoError(t, err)
 
 	tarFn := func(ctx context.Context, _, tok string) ([]*targets.Target, error) {
@@ -91,9 +90,3 @@ func (s *TestServer) AddResources(t *testing.T, p *authtokens.AuthToken, tars []
 	}
 	require.NoError(t, r.Refresh(ctx, cache.WithTargetRetrievalFunc(tarFn), cache.WithSessionRetrievalFunc(sessFn)))
 }
-
-// unimplementedAuthTokenReader is an unimplemented function for reading auth
-// tokens from a provided boundary address.
-func unimplementedAuthTokenReader(ctx context.Context, addr string, authToken string) (*authtokens.AuthToken, error) {
-	return nil, stdErrors.New("unimplemented")
-}
diff --git a/internal/cmd/commands/search/search_test.go b/internal/cmd/commands/search/search_test.go
index afc41d4862..541b07167c 100644
--- a/internal/cmd/commands/search/search_test.go
+++ b/internal/cmd/commands/search/search_test.go
@@ -56,18 +56,19 @@ func TestSearch(t *testing.T) {
 		ExpirationTime: time.Now().Add(time.Minute),
 	}
 	cmd := &testCommander{t: t, at: at}
+	boundaryTokenReaderFn := func(ctx context.Context, addr, authToken string) (*authtokens.AuthToken, error) {
+		if authToken == at.Token {
+			return at, nil
+		}
+		return nil, errors.New("test not found error")
+	}
 
 	srv := daemon.NewTestServer(t, cmd)
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		srv.Serve(t, daemon.WithBoundaryTokenReaderFunc(ctx, func(ctx context.Context, addr, authToken string) (*authtokens.AuthToken, error) {
-			if authToken == at.Token {
-				return at, nil
-			}
-			return nil, errors.New("test not found error")
-		}))
+		srv.Serve(t, daemon.WithBoundaryTokenReaderFunc(ctx, boundaryTokenReaderFn))
 	}()
 	// Give the store some time to get initialized
 	time.Sleep(100 * time.Millisecond)
@@ -161,7 +162,7 @@ func TestSearch(t *testing.T) {
 	}, []*sessions.Session{
 		{Id: "sess_1234567890", TargetId: "ttcp_1234567890", Status: "pending"},
 		{Id: "sess_0987654321", TargetId: "ttcp_0987654321", Status: "pending"},
-	})
+	}, boundaryTokenReaderFn)
 
 	t.Run("target response from list", func(t *testing.T) {
 		resp, err := search(ctx, srv.BaseSocketDir(), filterBy{