Merge both mocks

Author: simonswine

pkg/frontend/vcs/source/find_go_test.go

@@ -11,28 +11,14 @@ import (
 	"github.com/grafana/pyroscope/pkg/frontend/vcs/config"
 )
 
-type VCSClientMock struct {
-	fileToFind       string
-	searchedSequence []string
-}
-
-func (c *VCSClientMock) GetFile(ctx context.Context, req client.FileRequest) (client.File, error) {
-	c.searchedSequence = append(c.searchedSequence, req.Path)
-	if req.Path == c.fileToFind {
-		return client.File{}, nil
-	} else {
-		return client.File{}, client.ErrNotFound
-	}
-}
-
 func Test_tryFindGoFile(t *testing.T) {
 	pyroscopeRepo, _ := giturl.NewGitURL("http://github.com/grafana/pyroscope")
 	tests := []struct {
 		name                  string
 		searchedPath          string
 		rootPath              string
 		repo                  giturl.IGitURL
-		clientMock            *VCSClientMock
+		clientMock            *mockVCSClient
 		attempts              int
 		expectedSearchedPaths []string
 		expectedError         error
@@ -42,7 +28,7 @@ func Test_tryFindGoFile(t *testing.T) {
 			searchedPath:          "/var/service1/src/main.go",
 			rootPath:              "",
 			repo:                  pyroscopeRepo,
-			clientMock:            &VCSClientMock{fileToFind: "/main.go"},
+			clientMock:            newMockVCSClient().addFiles(newFile("main.go")),
 			attempts:              5,
 			expectedSearchedPaths: []string{"var/service1/src/main.go", "service1/src/main.go", "src/main.go", "main.go"},
 			expectedError:         nil,
@@ -52,7 +38,7 @@ func Test_tryFindGoFile(t *testing.T) {
 			searchedPath:          "/src/main.go",
 			rootPath:              "service/example",
 			repo:                  pyroscopeRepo,
-			clientMock:            &VCSClientMock{fileToFind: "service/example/main.go"},
+			clientMock:            newMockVCSClient().addFiles(newFile("service/example/main.go")),
 			attempts:              5,
 			expectedSearchedPaths: []string{"service/example/src/main.go", "service/example/main.go"},
 			expectedError:         nil,
@@ -62,7 +48,7 @@ func Test_tryFindGoFile(t *testing.T) {
 			searchedPath:          "github.com/grafana/pyroscope/main.go",
 			rootPath:              "",
 			repo:                  pyroscopeRepo,
-			clientMock:            &VCSClientMock{fileToFind: "/main.go"},
+			clientMock:            newMockVCSClient().addFiles(newFile("main.go")),
 			attempts:              1,
 			expectedSearchedPaths: []string{"main.go"},
 			expectedError:         nil,
@@ -72,7 +58,7 @@ func Test_tryFindGoFile(t *testing.T) {
 			searchedPath:          "/Users/pyroscope/git/github.com/grafana/pyroscope/main.go",
 			rootPath:              "",
 			repo:                  pyroscopeRepo,
-			clientMock:            &VCSClientMock{fileToFind: "/main.go"},
+			clientMock:            newMockVCSClient().addFiles(newFile("main.go")),
 			attempts:              1,
 			expectedSearchedPaths: []string{"main.go"},
 			expectedError:         nil,
@@ -82,7 +68,7 @@ func Test_tryFindGoFile(t *testing.T) {
 			searchedPath:          "/var/service1/src/main.go",
 			rootPath:              "",
 			repo:                  pyroscopeRepo,
-			clientMock:            &VCSClientMock{fileToFind: "/main.go"},
+			clientMock:            newMockVCSClient().addFiles(newFile("main.go")),
 			attempts:              3,
 			expectedSearchedPaths: []string{"var/service1/src/main.go", "service1/src/main.go", "src/main.go"},
 			expectedError:         client.ErrNotFound,
@@ -94,6 +80,7 @@ func Test_tryFindGoFile(t *testing.T) {
 			sut := FileFinder{
 				file:     config.FileSpec{Path: tt.searchedPath},
 				rootPath: tt.rootPath,
+				ref:      defaultRef(""),
 				repo:     tt.repo,
 				client:   tt.clientMock,
 			}

pkg/frontend/vcs/source/find_test.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"path/filepath"
+	"sync"
 	"testing"
 
 	"github.com/go-kit/log"
@@ -17,49 +18,66 @@ import (
 	"github.com/grafana/pyroscope/pkg/frontend/vcs/config"
 )
 
-func newTestVCSClient(b testing.TB) *testVCSClient {
-	return &testVCSClient{
+func newMockVCSClient() *mockVCSClient {
+	return &mockVCSClient{
 		files: make(map[client.FileRequest]client.File),
-		t:     b,
 	}
 }
 
-type testVCSClient struct {
-	files map[client.FileRequest]client.File
-	t     testing.TB
+type mockFileResponse struct {
+	request client.FileRequest
+	content string
 }
 
-func (m *testVCSClient) GetFile(ctx context.Context, req client.FileRequest) (client.File, error) {
-	if file, ok := m.files[req]; ok {
-		return file, nil
+func newFile(path string) mockFileResponse {
+	return mockFileResponse{
+		request: client.FileRequest{
+			Path: path,
+		},
+		content: "# Content of " + path,
 	}
-	m.t.Logf("file not found: %+v", req)
-	return client.File{}, client.ErrNotFound
 }
 
-// mockFileResponse represents a file to be mocked
-type mockFileResponse struct {
-	request client.FileRequest
-	content string
+func (f *mockFileResponse) url() string {
+	return fmt.Sprintf(
+		"https://github.com/%s/%s/blob/%s/%s",
+		f.request.Owner,
+		f.request.Repo,
+		f.request.Ref,
+		f.request.Path,
+	)
+}
+
+type mockVCSClient struct {
+	mtx              sync.Mutex
+	files            map[client.FileRequest]client.File
+	searchedSequence []string
+}
+
+func (c *mockVCSClient) GetFile(ctx context.Context, req client.FileRequest) (client.File, error) {
+	c.mtx.Lock()
+	c.searchedSequence = append(c.searchedSequence, req.Path)
+	file, ok := c.files[req]
+	c.mtx.Unlock()
+	if ok {
+		return file, nil
+	}
+	return client.File{}, client.ErrNotFound
 }
 
-// setupMockFiles is a shared function to setup mock responses for GetFile calls
-func (m *testVCSClient) addFiles(files ...mockFileResponse) {
+func (c *mockVCSClient) addFiles(files ...mockFileResponse) *mockVCSClient {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 	for _, file := range files {
 		file.request.Owner = defaultOwner(file.request.Owner)
 		file.request.Repo = defaultRepo(file.request.Repo)
 		file.request.Ref = defaultRef(file.request.Ref)
-		m.files[file.request] = client.File{
+		c.files[file.request] = client.File{
 			Content: file.content,
-			URL: fmt.Sprintf(
-				"https://github.com/%s/%s/blob/%s/%s",
-				file.request.Owner,
-				file.request.Repo,
-				file.request.Ref,
-				file.request.Path,
-			),
+			URL:     file.url(),
 		}
 	}
+	return c
 }
 
 func defaultOwner(s string) string {
@@ -393,7 +411,7 @@ require (
 			ctx := context.Background()
 
 			// Setup mock VCS client
-			mockClient := newTestVCSClient(t)
+			mockClient := newMockVCSClient()
 
 			// Populate pyroscopeYAML content into first mock file (if present)
 			mockFiles := append(tt.mockFiles, mockFileResponse{