From 284e3220408eb06b9388b0ef8cede6ade5c99cda Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 15:12:35 -0700 Subject: [PATCH 1/7] Ensure constructor checks for non 0 config values --- config/config.go | 12 ++++++++++++ controller/BUILD.bazel | 1 + controller/controller.go | 27 ++++++++++++--------------- controller/controller_test.go | 4 ++-- controller/getchangedtargets_test.go | 14 +++++++------- controller/gettargetgraph_test.go | 28 ++++++++++++++-------------- controller/testhelper_test.go | 20 ++++++++++++++++++++ core/bazel/bazel.go | 9 ++------- core/bazel/bazel_test.go | 4 ++++ core/bazel/query_test.go | 4 ++++ example/main.go | 6 +++++- example/tango-config.yaml | 4 ++++ orchestrator/testdata/config.yaml | 4 ++++ 13 files changed, 91 insertions(+), 46 deletions(-) diff --git a/config/config.go b/config/config.go index 77ca493..9facddb 100644 --- a/config/config.go +++ b/config/config.go @@ -70,12 +70,24 @@ func Parse(configFilePath string) (*Config, error) { if config.Service.WorkerPoolSize <= 0 { return nil, fmt.Errorf("service.worker_pool_size must be > 0, got %d", config.Service.WorkerPoolSize) } + if config.Service.Chunking.TargetChunkSize <= 0 { + return nil, fmt.Errorf("service.chunking.target_chunk_size must be > 0, got %d", config.Service.Chunking.TargetChunkSize) + } + if config.Service.Chunking.ChangedTargetChunkSize <= 0 { + return nil, fmt.Errorf("service.chunking.changed_target_chunk_size must be > 0, got %d", config.Service.Chunking.ChangedTargetChunkSize) + } + if config.Service.Chunking.MetadataMapChunkSize <= 0 { + return nil, fmt.Errorf("service.chunking.metadata_map_chunk_size must be > 0, got %d", config.Service.Chunking.MetadataMapChunkSize) + } config.repositoryByRemote = make(map[string]*RepositoryConfig, len(config.Repository)) for i := range config.Repository { remote := config.Repository[i].Remote if remote == "" { return nil, fmt.Errorf("repository[%d].remote must not be empty", i) } + if config.Repository[i].QueryTimeout <= 0 { + return nil, fmt.Errorf("repository[%d].query_timeout must be > 0", i) + } if _, exists := config.repositoryByRemote[remote]; exists { return nil, fmt.Errorf("duplicate repository remote %q", remote) } diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 581ba3b..107d58e 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -37,6 +37,7 @@ go_test( ], embed = [":controller"], deps = [ + "//config", "//core/common", "//core/storage", "//core/storage/storagemock", diff --git a/controller/controller.go b/controller/controller.go index 34ca33c..92089c9 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -16,11 +16,11 @@ package controller import ( "context" + "fmt" "time" "github.com/uber-go/tally" "github.com/uber/tango/config" - "github.com/uber/tango/core/common" "github.com/uber/tango/core/storage" "github.com/uber/tango/orchestrator" pb "github.com/uber/tango/tangopb" @@ -59,34 +59,31 @@ type controller struct { // NewController creates a new controller. appCtx is cancelled on process // shutdown to abort background work. -func NewController(appCtx context.Context, p Params) pb.TangoYARPCServer { +func NewController(appCtx context.Context, p Params) (pb.TangoYARPCServer, error) { scope := p.Scope if scope == nil { scope = tally.NoopScope } - targetChunkSize := p.ChunkConfig.TargetChunkSize - if targetChunkSize <= 0 { - targetChunkSize = common.DefaultTargetChunkSize + if p.ChunkConfig.TargetChunkSize <= 0 { + return nil, fmt.Errorf("target_chunk_size must be > 0, got %d", p.ChunkConfig.TargetChunkSize) } - changedTargetChunkSize := p.ChunkConfig.ChangedTargetChunkSize - if changedTargetChunkSize <= 0 { - changedTargetChunkSize = common.DefaultChangedTargetChunkSize + if p.ChunkConfig.ChangedTargetChunkSize <= 0 { + return nil, fmt.Errorf("changed_target_chunk_size must be > 0, got %d", p.ChunkConfig.ChangedTargetChunkSize) } - metadataMapChunkSize := p.ChunkConfig.MetadataMapChunkSize - if metadataMapChunkSize <= 0 { - metadataMapChunkSize = common.DefaultMetadataMapChunkSize + if p.ChunkConfig.MetadataMapChunkSize <= 0 { + return nil, fmt.Errorf("metadata_map_chunk_size must be > 0, got %d", p.ChunkConfig.MetadataMapChunkSize) } return &controller{ logger: p.Logger, storage: p.Storage, orchestrator: p.Orchestrator, scope: scope.SubScope("controller"), - targetChunkSize: targetChunkSize, - changedTargetChunkSize: changedTargetChunkSize, - metadataMapChunkSize: metadataMapChunkSize, + targetChunkSize: p.ChunkConfig.TargetChunkSize, + changedTargetChunkSize: p.ChunkConfig.ChangedTargetChunkSize, + metadataMapChunkSize: p.ChunkConfig.MetadataMapChunkSize, totalDurationBuckets: _totalDurationBuckets, appCtx: appCtx, - } + }, nil } // linkRequestCtx returns a context derived from reqCtx that is also cancelled diff --git a/controller/controller_test.go b/controller/controller_test.go index b891f87..a4f297e 100644 --- a/controller/controller_test.go +++ b/controller/controller_test.go @@ -31,10 +31,10 @@ func TestNewController_StoresAppContext(t *testing.T) { appCtx, cancel := context.WithCancel(context.Background()) defer cancel() - c := NewController(appCtx, Params{ + c := mustNewController(t, appCtx, Params{ Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), - }).(*controller) + }) assert.Same(t, appCtx, c.appCtx) assert.NoError(t, c.appCtx.Err()) diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index 1e854fc..65d3cff 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -155,7 +155,7 @@ func TestGetChangedTargets_ValidationError(t *testing.T) { ctrl := gomock.NewController(t) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) - c := NewController(context.Background(), Params{Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl)}) + c := mustNewController(t, context.Background(), Params{Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl)}) err := c.GetChangedTargets(nil, stream) assert.EqualError(t, err, "request cannot be nil") @@ -196,7 +196,7 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { stream.EXPECT().Send(gomock.Any()).Return(nil).Times(2) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), @@ -227,7 +227,7 @@ func TestGetChangedTargets_TreehashReadError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{}, injected).Times(2) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zap.NewNop(), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), @@ -316,7 +316,7 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { return nil }) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), @@ -423,7 +423,7 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { return nil }) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), @@ -523,7 +523,7 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { appCtx, cancelApp := context.WithCancel(context.Background()) defer cancelApp() - c := NewController(appCtx, Params{ + c := mustNewController(t, appCtx, Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), @@ -1159,7 +1159,7 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { return nil }).Times(2) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index e62856f..9b6f6cf 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -46,7 +46,7 @@ func TestGetTargetGraph_CacheMiss_NoSend(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte{})}, nil), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -71,7 +71,7 @@ func TestGetTargetGraph_StorageError_Propagates(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) storagemock := storagemock.NewMockStorage(ctrl) storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, expected) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, }) @@ -97,7 +97,7 @@ func TestGetTargetGraph_DecodeError_ReturnsError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("bad-bytes"))}, nil), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, }) @@ -130,7 +130,7 @@ func TestGetTargetGraph_SendsWhenItemPresent(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-xyz"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -152,7 +152,7 @@ func TestGetTargetGraph_BuildDescriptionMissingRequiredFields_ReturnsError(t *te stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -173,7 +173,7 @@ func TestGetTargetGraph_MissingBuildDescription_ReturnsError(t *testing.T) { stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -202,7 +202,7 @@ func TestGetTargetGraph_TreehashNotFound_NoError(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orchestrator, @@ -220,7 +220,7 @@ func TestGetTargetGraph_TreehashReadError(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -240,7 +240,7 @@ func TestGetTargetGraph_GraphFetchError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("graph error")), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -264,7 +264,7 @@ func TestGetTargetGraph_GraphReadError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -292,7 +292,7 @@ func TestGetTargetGraph_StreamSendError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, }) @@ -321,7 +321,7 @@ func TestGetTargetGraph_GraphNotFound_FallsThrough(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, @@ -343,7 +343,7 @@ func TestGetTargetGraph_GraphReadCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("context canceled")), ) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, }) @@ -367,7 +367,7 @@ func TestGetTargetGraph_OrchestratorCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, &storage.NotFoundError{Path: "x"}) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(nil, errors.New("context canceled")) - c := NewController(context.Background(), Params{ + c := mustNewController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go index 56ff439..fc87b42 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -16,12 +16,20 @@ package controller import ( "context" + "testing" "github.com/uber-go/tally" + "github.com/uber/tango/config" "github.com/uber/tango/core/common" "go.uber.org/zap" ) +var _testChunkConfig = config.ChunkConfig{ + TargetChunkSize: common.DefaultTargetChunkSize, + ChangedTargetChunkSize: common.DefaultChangedTargetChunkSize, + MetadataMapChunkSize: common.DefaultMetadataMapChunkSize, +} + func newTestController(logger *zap.Logger) *controller { return &controller{ logger: logger, @@ -33,3 +41,15 @@ func newTestController(logger *zap.Logger) *controller { appCtx: context.Background(), } } + +func mustNewController(t *testing.T, appCtx context.Context, p Params) *controller { + t.Helper() + if p.ChunkConfig == (config.ChunkConfig{}) { + p.ChunkConfig = _testChunkConfig + } + c, err := NewController(appCtx, p) + if err != nil { + t.Fatalf("NewController: %v", err) + } + return c.(*controller) +} diff --git a/core/bazel/bazel.go b/core/bazel/bazel.go index 217283a..033930e 100644 --- a/core/bazel/bazel.go +++ b/core/bazel/bazel.go @@ -30,11 +30,6 @@ import ( "go.uber.org/zap" ) -const ( - // default query timeout if not provided in config - _queryTimeout = 15 * time.Minute -) - type QueryRequest struct { Query string StartupOptions []string @@ -84,8 +79,8 @@ func NewBazelClient(ctx context.Context, p Params) (*BazelClient, error) { } } timeout := p.QueryTimeout - if timeout == 0 { - timeout = _queryTimeout + if timeout <= 0 { + return nil, fmt.Errorf("query timeout must be > 0, got %s", timeout) } bazelCommand, err := detectBazelExecutable(ctx, p.BazelCommand) if err != nil { diff --git a/core/bazel/bazel_test.go b/core/bazel/bazel_test.go index c76342a..409e918 100644 --- a/core/bazel/bazel_test.go +++ b/core/bazel/bazel_test.go @@ -18,6 +18,7 @@ import ( "context" "os/exec" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -36,6 +37,7 @@ func TestNewBazelClient(t *testing.T) { WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{"FOO": "bar"}, Logger: zap.NewNop().Sugar(), + QueryTimeout: 5 * time.Minute, ExecCommandContext: func(ctx context.Context, name string, arg ...string) commander { return nil }, @@ -48,6 +50,7 @@ func TestNewBazelClient(t *testing.T) { WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{"FOO": "bar"}, Logger: zap.NewNop().Sugar(), + QueryTimeout: 5 * time.Minute, }, }, } @@ -72,6 +75,7 @@ func TestNewBazelClient_WithNilExecCommand(t *testing.T) { WorkspacePath: "/workspace", EnvVarsMap: map[string]string{"KEY": "value"}, Logger: zap.NewNop().Sugar(), + QueryTimeout: 5 * time.Minute, }) require.NoError(t, err) require.NotNil(t, client) diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index c506214..2b72982 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -57,6 +57,7 @@ func TestExecuteQuery_Success(t *testing.T) { mockCmd.EXPECT().Wait().Return(nil), ) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -101,6 +102,7 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { mockCmd.EXPECT().Wait().Return(nil), ) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -230,6 +232,7 @@ func TestExecuteQueryInternal_Failures(t *testing.T) { tt.setupMock(mockCmd) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -259,6 +262,7 @@ func TestExecuteQuery_ErrorCase(t *testing.T) { mockCmd.EXPECT().StdoutPipe().Return(nil, errors.New("stdout pipe failed")) client, err := NewBazelClient(context.Background(), Params{ + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, diff --git a/example/main.go b/example/main.go index 05a3036..c5c900d 100644 --- a/example/main.go +++ b/example/main.go @@ -105,11 +105,15 @@ func run() error { // Controller (YARPC server implementation). appCtx is forwarded so the // controller's background goroutines are tied to process lifetime. - ctrl := controller.NewController(appCtx, controller.Params{ + ctrl, err := controller.NewController(appCtx, controller.Params{ Logger: zl, Storage: store, Orchestrator: orch, + ChunkConfig: cfg.Service.Chunking, }) + if err != nil { + return fmt.Errorf("failed to create controller: %w", err) + } // YARPC transports and dispatcher grpcTransport := yarpcgrpc.NewTransport() diff --git a/example/tango-config.yaml b/example/tango-config.yaml index ec43e6a..1e45c0d 100644 --- a/example/tango-config.yaml +++ b/example/tango-config.yaml @@ -24,3 +24,7 @@ service: repo_manager_clone_path: "/tmp/tango-repo-manager" # root for worker checkouts; defaults to repo_manager_clone_path/.workers worker_root_path: "/tmp/tango/workers" + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 diff --git a/orchestrator/testdata/config.yaml b/orchestrator/testdata/config.yaml index f420f74..34887e0 100644 --- a/orchestrator/testdata/config.yaml +++ b/orchestrator/testdata/config.yaml @@ -11,3 +11,7 @@ repository: service: worker_pool_size: 3 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 From 1616786ac111492f255467136a999b7f506d9b88 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 15:52:19 -0700 Subject: [PATCH 2/7] Update constructors --- controller/controller_test.go | 6 +- controller/getchangedtargets_test.go | 43 +++++++--- controller/gettargetgraph_test.go | 124 ++++++++++++++++----------- controller/testhelper_test.go | 13 --- 4 files changed, 111 insertions(+), 75 deletions(-) diff --git a/controller/controller_test.go b/controller/controller_test.go index a4f297e..3aa7385 100644 --- a/controller/controller_test.go +++ b/controller/controller_test.go @@ -19,6 +19,7 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" orchestratormock "github.com/uber/tango/orchestrator/orchestratormock" "go.uber.org/mock/gomock" "go.uber.org/zap" @@ -31,10 +32,13 @@ func TestNewController_StoresAppContext(t *testing.T) { appCtx, cancel := context.WithCancel(context.Background()) defer cancel() - c := mustNewController(t, appCtx, Params{ + server, err := NewController(appCtx, Params{ Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) + c := server.(*controller) assert.Same(t, appCtx, c.appCtx) assert.NoError(t, c.appCtx.Err()) diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index 65d3cff..c991039 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -155,9 +155,14 @@ func TestGetChangedTargets_ValidationError(t *testing.T) { ctrl := gomock.NewController(t) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) - c := mustNewController(t, context.Background(), Params{Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl)}) + c, err := NewController(context.Background(), Params{ + Logger: zap.NewNop(), + Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, + }) + require.NoError(t, err) - err := c.GetChangedTargets(nil, stream) + err = c.GetChangedTargets(nil, stream) assert.EqualError(t, err, "request cannot be nil") } @@ -196,11 +201,13 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { stream.EXPECT().Send(gomock.Any()).Return(nil).Times(2) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -208,7 +215,7 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1}, } - err := c.GetChangedTargets(request, stream) + err = c.GetChangedTargets(request, stream) require.NoError(t, err) } @@ -227,11 +234,13 @@ func TestGetChangedTargets_TreehashReadError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{}, injected).Times(2) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zap.NewNop(), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -239,7 +248,7 @@ func TestGetChangedTargets_TreehashReadError(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1}, } - err := c.GetChangedTargets(request, stream) + err = c.GetChangedTargets(request, stream) require.Error(t, err) require.ErrorIs(t, err, injected) var ce common.ClassifiedError @@ -316,11 +325,13 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { return nil }) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -328,7 +339,7 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1}, } - err := c.GetChangedTargets(request, stream) + err = c.GetChangedTargets(request, stream) assert.Error(t, err) select { @@ -423,11 +434,13 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { return nil }) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -435,7 +448,7 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1, IncludeHashes: true, IncludeTags: true, IncludeAttributes: true}, } - err := c.GetChangedTargets(request, stream) + err = c.GetChangedTargets(request, stream) require.NoError(t, err) select { @@ -523,11 +536,13 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { appCtx, cancelApp := context.WithCancel(context.Background()) defer cancelApp() - c := mustNewController(t, appCtx, Params{ + c, err := NewController(appCtx, Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -1159,11 +1174,13 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { return nil }).Times(2) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -1171,7 +1188,7 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: 1}, } - err := c.GetChangedTargets(request, stream) + err = c.GetChangedTargets(request, stream) require.NoError(t, err) require.Len(t, sent, 2) diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index 9b6f6cf..23277b6 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -46,10 +46,12 @@ func TestGetTargetGraph_CacheMiss_NoSend(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte{})}, nil), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) req := &pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", @@ -60,7 +62,7 @@ func TestGetTargetGraph_CacheMiss_NoSend(t *testing.T) { }, }, } - err := c.GetTargetGraph(req, stream) + err = c.GetTargetGraph(req, stream) require.NoError(t, err) } @@ -71,11 +73,13 @@ func TestGetTargetGraph_StorageError_Propagates(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) storagemock := storagemock.NewMockStorage(ctrl) storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, expected) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: storagemock, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: storagemock, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", BaseSha: "sha", @@ -97,11 +101,13 @@ func TestGetTargetGraph_DecodeError_ReturnsError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("bad-bytes"))}, nil), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: storagemock, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: storagemock, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", BaseSha: "sha", @@ -130,10 +136,12 @@ func TestGetTargetGraph_SendsWhenItemPresent(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-xyz"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", @@ -152,11 +160,13 @@ func TestGetTargetGraph_BuildDescriptionMissingRequiredFields_ReturnsError(t *te stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", Requests: []*pb.Request{ @@ -173,11 +183,13 @@ func TestGetTargetGraph_MissingBuildDescription_ReturnsError(t *testing.T) { stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{}, stream) + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{}, stream) assert.Error(t, err) } @@ -202,12 +214,14 @@ func TestGetTargetGraph_TreehashNotFound_NoError(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orchestrator, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.NoError(t, err) @@ -220,11 +234,13 @@ func TestGetTargetGraph_TreehashReadError(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) assert.Error(t, err) @@ -240,11 +256,13 @@ func TestGetTargetGraph_GraphFetchError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("graph error")), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.Error(t, err) @@ -264,11 +282,13 @@ func TestGetTargetGraph_GraphReadError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) assert.Error(t, err) @@ -292,10 +312,12 @@ func TestGetTargetGraph_StreamSendError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: storagemock, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: storagemock, + ChunkConfig: _testChunkConfig, }) + require.NoError(t, err) err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) @@ -321,12 +343,14 @@ func TestGetTargetGraph_GraphNotFound_FallsThrough(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.NoError(t, err) @@ -343,11 +367,13 @@ func TestGetTargetGraph_GraphReadCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("context canceled")), ) - c := mustNewController(t, context.Background(), Params{ - Logger: zaptest.NewLogger(t), - Storage: store, + c, err := NewController(context.Background(), Params{ + Logger: zaptest.NewLogger(t), + Storage: store, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.Error(t, err) @@ -367,12 +393,14 @@ func TestGetTargetGraph_OrchestratorCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, &storage.NotFoundError{Path: "x"}) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(nil, errors.New("context canceled")) - c := mustNewController(t, context.Background(), Params{ + c, err := NewController(context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, + ChunkConfig: _testChunkConfig, }) - err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ + require.NoError(t, err) + err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.Error(t, err) diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go index fc87b42..e01f23f 100644 --- a/controller/testhelper_test.go +++ b/controller/testhelper_test.go @@ -16,7 +16,6 @@ package controller import ( "context" - "testing" "github.com/uber-go/tally" "github.com/uber/tango/config" @@ -41,15 +40,3 @@ func newTestController(logger *zap.Logger) *controller { appCtx: context.Background(), } } - -func mustNewController(t *testing.T, appCtx context.Context, p Params) *controller { - t.Helper() - if p.ChunkConfig == (config.ChunkConfig{}) { - p.ChunkConfig = _testChunkConfig - } - c, err := NewController(appCtx, p) - if err != nil { - t.Fatalf("NewController: %v", err) - } - return c.(*controller) -} From b127d12b5597ac8ff35473f3b6a17d887c7547fa Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 16:44:13 -0700 Subject: [PATCH 3/7] Update --- controller/BUILD.bazel | 2 - controller/controller_test.go | 18 ++++++-- controller/getchangedtargets_test.go | 53 ++++++++++------------ controller/gettargetgraph_test.go | 66 +++++++++++----------------- controller/testhelper_test.go | 42 ------------------ core/common/utils.go | 36 +++------------ core/common/utils_test.go | 6 +-- example/cmd/query-bench/main.go | 2 +- orchestrator/native_orchestrator.go | 2 +- 9 files changed, 75 insertions(+), 152 deletions(-) delete mode 100644 controller/testhelper_test.go diff --git a/controller/BUILD.bazel b/controller/BUILD.bazel index 107d58e..2194182 100644 --- a/controller/BUILD.bazel +++ b/controller/BUILD.bazel @@ -33,7 +33,6 @@ go_test( "getchangedtargets_test.go", "gettargetgraph_test.go", "output_filter_test.go", - "testhelper_test.go", ], embed = [":controller"], deps = [ @@ -47,7 +46,6 @@ go_test( "@com_github_gogo_protobuf//io", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", - "@com_github_uber_go_tally//:tally", "@org_uber_go_mock//gomock", "@org_uber_go_zap//:zap", "@org_uber_go_zap//zaptest", diff --git a/controller/controller_test.go b/controller/controller_test.go index 3aa7385..870b5c6 100644 --- a/controller/controller_test.go +++ b/controller/controller_test.go @@ -20,11 +20,25 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/uber/tango/config" orchestratormock "github.com/uber/tango/orchestrator/orchestratormock" "go.uber.org/mock/gomock" "go.uber.org/zap" ) +var _testChunkConfig = config.ChunkConfig{ + TargetChunkSize: 250, + ChangedTargetChunkSize: 125, + MetadataMapChunkSize: 50_000, +} + +func newTestController(t *testing.T, appCtx context.Context, p Params) *controller { + t.Helper() + ctrl, err := NewController(appCtx, p) + require.NoError(t, err) + return ctrl.(*controller) +} + // TestNewController_StoresAppContext verifies the caller-supplied context is // retained and is the one observed by background goroutines. func TestNewController_StoresAppContext(t *testing.T) { @@ -32,13 +46,11 @@ func TestNewController_StoresAppContext(t *testing.T) { appCtx, cancel := context.WithCancel(context.Background()) defer cancel() - server, err := NewController(appCtx, Params{ + c := newTestController(t, appCtx, Params{ Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - c := server.(*controller) assert.Same(t, appCtx, c.appCtx) assert.NoError(t, c.appCtx.Err()) diff --git a/controller/getchangedtargets_test.go b/controller/getchangedtargets_test.go index c991039..02dda25 100644 --- a/controller/getchangedtargets_test.go +++ b/controller/getchangedtargets_test.go @@ -133,7 +133,7 @@ func TestValidateGetChangedTargetsRequest(t *testing.T) { } func TestCompareTargetGraphs(t *testing.T) { - c := newTestController(zap.NewNop()) + c := newTestController(t, context.Background(), Params{Logger: zap.NewNop(), ChunkConfig: _testChunkConfig}) firstGraph := &pb.GetTargetGraphResponse{ Item: &pb.GetTargetGraphResponse_Metadata{ @@ -155,14 +155,13 @@ func TestGetChangedTargets_ValidationError(t *testing.T) { ctrl := gomock.NewController(t) stream := tangomock.NewMockTangoServiceGetChangedTargetsYARPCServer(ctrl) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zap.NewNop(), Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetChangedTargets(nil, stream) + err := c.GetChangedTargets(nil, stream) assert.EqualError(t, err, "request cannot be nil") } @@ -201,13 +200,12 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { stream.EXPECT().Send(gomock.Any()).Return(nil).Times(2) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -215,7 +213,7 @@ func TestGetChangedTargets_CacheHit(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1}, } - err = c.GetChangedTargets(request, stream) + err := c.GetChangedTargets(request, stream) require.NoError(t, err) } @@ -234,13 +232,12 @@ func TestGetChangedTargets_TreehashReadError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{}, injected).Times(2) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zap.NewNop(), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -248,7 +245,7 @@ func TestGetChangedTargets_TreehashReadError(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1}, } - err = c.GetChangedTargets(request, stream) + err := c.GetChangedTargets(request, stream) require.Error(t, err) require.ErrorIs(t, err, injected) var ce common.ClassifiedError @@ -325,13 +322,12 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { return nil }) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -339,7 +335,7 @@ func TestGetChangedTargets_StreamSendError(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1}, } - err = c.GetChangedTargets(request, stream) + err := c.GetChangedTargets(request, stream) assert.Error(t, err) select { @@ -434,13 +430,12 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { return nil }) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -448,7 +443,7 @@ func TestGetChangedTargets_streamChunks(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: -1, IncludeHashes: true, IncludeTags: true, IncludeAttributes: true}, } - err = c.GetChangedTargets(request, stream) + err := c.GetChangedTargets(request, stream) require.NoError(t, err) select { @@ -536,13 +531,12 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { appCtx, cancelApp := context.WithCancel(context.Background()) defer cancelApp() - c, err := NewController(appCtx, Params{ + c := newTestController(t, appCtx, Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -591,7 +585,7 @@ func TestGetChangedTargets_CacheWriteUsesAppCtx(t *testing.T) { } func TestCompareTargetGraphs_NewTarget_CanonicalIDs(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) first := []*pb.GetTargetGraphResponse{ { @@ -644,7 +638,7 @@ func TestCompareTargetGraphs_NewTarget_CanonicalIDs(t *testing.T) { } func TestCompareTargetGraphs_SourceFileDirectAndPropagation(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: source file A (id 1, hash h1), lib L (id 2, hash h1, dep -> A) first := []*pb.GetTargetGraphResponse{ @@ -728,7 +722,7 @@ func TestCompareTargetGraphs_SourceFileDirectAndPropagation(t *testing.T) { } func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (id 1, rule), no deps first := []*pb.GetTargetGraphResponse{ @@ -784,7 +778,7 @@ func TestCompareTargetGraphs_ChangedRuleUnreachableFromAnySeed(t *testing.T) { } func TestCompareTargetGraphs_ChangedWhenDependenciesChanged(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (id 1, rule) with deps on A first := []*pb.GetTargetGraphResponse{ @@ -860,7 +854,7 @@ func TestCompareTargetGraphs_ChangedWhenDependenciesChanged(t *testing.T) { } func TestCompareTargetGraphs_ChangedWhenAttributesChanged(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T with attribute "key1" -> "value1" first := []*pb.GetTargetGraphResponse{ @@ -933,7 +927,7 @@ func TestCompareTargetGraphs_ChangedWhenAttributesChanged(t *testing.T) { } func TestCompareTargetGraphs_ChangedWhenNewAttributeAdded(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T with one attribute first := []*pb.GetTargetGraphResponse{ @@ -1174,13 +1168,12 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { return nil }).Times(2) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, Orchestrator: orchestratormock.NewMockOrchestrator(ctrl), ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) request := &pb.GetChangedTargetsRequest{ FirstRevision: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha1"}, @@ -1188,7 +1181,7 @@ func TestGetChangedTargets_CacheHitWithDistanceFilter(t *testing.T) { OutputConfig: &pb.OutputConfig{MaxDistance: 1}, } - err = c.GetChangedTargets(request, stream) + err := c.GetChangedTargets(request, stream) require.NoError(t, err) require.Len(t, sent, 2) @@ -1273,7 +1266,7 @@ func TestComputeDistances_NilMetadata(t *testing.T) { } func TestCompareTargetGraphs_HashOnlyChangePropagatesViaBFS(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (rule) with deps on source file A (id 10) and attributes first := []*pb.GetTargetGraphResponse{ @@ -1366,7 +1359,7 @@ func TestCompareTargetGraphs_HashOnlyChangePropagatesViaBFS(t *testing.T) { } func TestCompareTargetGraphs_SiblingRuleNotPromotedToSeed(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Source file A (id 1) owned by rule L (id 2). // Rule T (id 3) depends on L (sibling rule), NOT directly on A. @@ -1445,7 +1438,7 @@ func TestCompareTargetGraphs_SiblingRuleNotPromotedToSeed(t *testing.T) { } func TestCompareTargetGraphs_DeletedTargetEmitted(t *testing.T) { - c := newTestController(zaptest.NewLogger(t)) + c := newTestController(t, context.Background(), Params{Logger: zaptest.NewLogger(t), ChunkConfig: _testChunkConfig}) // Old: T (rule) exists; New: T is gone. first := []*pb.GetTargetGraphResponse{ diff --git a/controller/gettargetgraph_test.go b/controller/gettargetgraph_test.go index 23277b6..ed55bd8 100644 --- a/controller/gettargetgraph_test.go +++ b/controller/gettargetgraph_test.go @@ -46,12 +46,11 @@ func TestGetTargetGraph_CacheMiss_NoSend(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()). Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte{})}, nil), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) req := &pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", @@ -62,7 +61,7 @@ func TestGetTargetGraph_CacheMiss_NoSend(t *testing.T) { }, }, } - err = c.GetTargetGraph(req, stream) + err := c.GetTargetGraph(req, stream) require.NoError(t, err) } @@ -73,13 +72,12 @@ func TestGetTargetGraph_StorageError_Propagates(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) storagemock := storagemock.NewMockStorage(ctrl) storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, expected) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", BaseSha: "sha", @@ -101,13 +99,12 @@ func TestGetTargetGraph_DecodeError_ReturnsError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("bad-bytes"))}, nil), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", BaseSha: "sha", @@ -136,12 +133,11 @@ func TestGetTargetGraph_SendsWhenItemPresent(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-xyz"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", @@ -160,13 +156,12 @@ func TestGetTargetGraph_BuildDescriptionMissingRequiredFields_ReturnsError(t *te stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{ Remote: "repo:go-code", Requests: []*pb.Request{ @@ -183,13 +178,12 @@ func TestGetTargetGraph_MissingBuildDescription_ReturnsError(t *testing.T) { stream := tangomock.NewMockTangoServiceGetTargetGraphYARPCServer(ctrl) stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{}, stream) + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{}, stream) assert.Error(t, err) } @@ -214,14 +208,13 @@ func TestGetTargetGraph_TreehashNotFound_NoError(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orchestrator.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orchestrator, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.NoError(t, err) @@ -234,13 +227,12 @@ func TestGetTargetGraph_TreehashReadError(t *testing.T) { stream.EXPECT().Context().Return(context.Background()) store := storagemock.NewMockStorage(ctrl) store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) assert.Error(t, err) @@ -256,13 +248,12 @@ func TestGetTargetGraph_GraphFetchError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("graph error")), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.Error(t, err) @@ -282,13 +273,12 @@ func TestGetTargetGraph_GraphReadError(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: &errReadCloser{err: errors.New("readfail")}}, nil), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) assert.Error(t, err) @@ -312,12 +302,11 @@ func TestGetTargetGraph_StreamSendError(t *testing.T) { storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), storagemock.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser(buf.Bytes())}, nil), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: storagemock, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) @@ -343,14 +332,13 @@ func TestGetTargetGraph_GraphNotFound_FallsThrough(t *testing.T) { graphReader.EXPECT().Read().Return(nil, io.EOF).Times(1) graphReader.EXPECT().Close().Return(nil) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(graphReader, nil) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.NoError(t, err) @@ -367,13 +355,12 @@ func TestGetTargetGraph_GraphReadCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{ReadCloser: newMockReadCloser([]byte("treehash-abc"))}, nil), store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, errors.New("context canceled")), ) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.Error(t, err) @@ -393,14 +380,13 @@ func TestGetTargetGraph_OrchestratorCancelled(t *testing.T) { store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, &storage.NotFoundError{Path: "x"}) orch := orchestratormock.NewMockOrchestrator(ctrl) orch.EXPECT().GetTargetGraph(gomock.Any(), gomock.Any()).Return(nil, errors.New("context canceled")) - c, err := NewController(context.Background(), Params{ + c := newTestController(t, context.Background(), Params{ Logger: zaptest.NewLogger(t), Storage: store, Orchestrator: orch, ChunkConfig: _testChunkConfig, }) - require.NoError(t, err) - err = c.GetTargetGraph(&pb.GetTargetGraphRequest{ + err := c.GetTargetGraph(&pb.GetTargetGraphRequest{ BuildDescription: &pb.BuildDescription{Remote: "repo:go-code", BaseSha: "sha"}, }, stream) require.Error(t, err) diff --git a/controller/testhelper_test.go b/controller/testhelper_test.go deleted file mode 100644 index e01f23f..0000000 --- a/controller/testhelper_test.go +++ /dev/null @@ -1,42 +0,0 @@ -// Copyright (c) 2026 Uber Technologies, Inc. -// -// 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 controller - -import ( - "context" - - "github.com/uber-go/tally" - "github.com/uber/tango/config" - "github.com/uber/tango/core/common" - "go.uber.org/zap" -) - -var _testChunkConfig = config.ChunkConfig{ - TargetChunkSize: common.DefaultTargetChunkSize, - ChangedTargetChunkSize: common.DefaultChangedTargetChunkSize, - MetadataMapChunkSize: common.DefaultMetadataMapChunkSize, -} - -func newTestController(logger *zap.Logger) *controller { - return &controller{ - logger: logger, - scope: tally.NoopScope, - targetChunkSize: common.DefaultTargetChunkSize, - changedTargetChunkSize: common.DefaultChangedTargetChunkSize, - metadataMapChunkSize: common.DefaultMetadataMapChunkSize, - totalDurationBuckets: _totalDurationBuckets, - appCtx: context.Background(), - } -} diff --git a/core/common/utils.go b/core/common/utils.go index 90df802..e46a170 100644 --- a/core/common/utils.go +++ b/core/common/utils.go @@ -28,24 +28,6 @@ import ( "github.com/uber/tango/tangopb" ) -const ( - // DefaultTargetChunkSize is the default number of OptimizedTarget entries per stream message. - // Sized conservatively: at ~40KB/target worst-case (target with ~10K direct deps × 4 bytes), - // 250 targets ≈ 10MB — well under the 64MB default gRPC per-message limit. - DefaultTargetChunkSize = 250 - - // DefaultChangedTargetChunkSize is the default number of ChangedTarget entries per stream message. - // A ChangedTarget carries both old_target and new_target (2× an OptimizedTarget), so we use - // half the regular chunk size to stay within the same byte budget. - DefaultChangedTargetChunkSize = 125 - - // DefaultMetadataMapChunkSize is the max entries per metadata message chunk. - // target_id_mapping and attribute_string_value_mapping scale with repo size and can exceed - // the 64MB gRPC message limit for large monorepos, so they are split across multiple messages. - // At ~85 bytes/entry (60-char avg target name + proto overhead), 50 000 entries ≈ 4.25MB per chunk. - DefaultMetadataMapChunkSize = 50_000 -) - // ToShortRemote returns the short remote name given a git ssh remote string. // For example, "git@github:uber/tango" will return "uber/tango". func ToShortRemote(remote string) string { @@ -137,8 +119,10 @@ func HashRequestOptions(opts *tangopb.RequestOptions) string { // for typical target rates. const cancelCheckInterval = 4096 -// ResultToGetTargetGraphResponse converts a Result to a GetTargetGraphResponse -func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result) ([]*tangopb.GetTargetGraphResponse, error) { +// ResultToGetTargetGraphResponse converts a Result to a GetTargetGraphResponse. +// targetChunkSize controls how many OptimizedTarget entries per stream message. +// metadataMapChunkSize controls how many entries per metadata map chunk. +func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Result, targetChunkSize, metadataMapChunkSize int) ([]*tangopb.GetTargetGraphResponse, error) { // Map target names to ids. This list is topologically sorted, so the ids are stable. // IDs start at 1 — 0 is reserved as the proto3 "unset" sentinel so consumers using // encoding/json (which honors `omitempty` on int32 fields) never silently lose a target. @@ -224,14 +208,14 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res attrStrValIDToVal := attrStrValMapper.Invert() // chunk targets into multiple messages for streaming - responses := chunkTargets(optimizedTargets, DefaultTargetChunkSize) + responses := chunkTargets(optimizedTargets, targetChunkSize) for _, meta := range ChunkMetadata( targetIDToName, ruleTypeIDToName, tagIDToName, attrNameIDToName, attrStrValIDToVal, - DefaultMetadataMapChunkSize, + metadataMapChunkSize, ) { responses = append(responses, &tangopb.GetTargetGraphResponse{ Item: &tangopb.GetTargetGraphResponse_Metadata{Metadata: meta}, @@ -242,10 +226,6 @@ func ResultToGetTargetGraphResponse(ctx context.Context, result targethasher.Res } func chunkTargets(targets []*tangopb.OptimizedTarget, chunkSize int) []*tangopb.GetTargetGraphResponse { - if chunkSize <= 0 { - chunkSize = DefaultTargetChunkSize - } - // at least one chunk numChunks := max(1, (len(targets)+chunkSize-1)/chunkSize) @@ -293,10 +273,6 @@ func ChunkMetadata( attrStrValIDToVal map[int32]string, chunkSize int, ) []*tangopb.Metadata { - if chunkSize <= 0 { - chunkSize = DefaultMetadataMapChunkSize - } - targetChunks := splitMap(targetIDToName, chunkSize) attrValChunks := splitMap(attrStrValIDToVal, chunkSize) diff --git a/core/common/utils_test.go b/core/common/utils_test.go index 4172da0..ea1ecd9 100644 --- a/core/common/utils_test.go +++ b/core/common/utils_test.go @@ -186,7 +186,7 @@ func TestChunkTargets(t *testing.T) { func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { t.Parallel() - // 500 targets with DefaultTargetChunkSize=250 → 2 target chunks + 1 metadata = 3 responses + // 500 targets with targetChunkSize=250 → 2 target chunks + 1 metadata = 3 responses numTargets := 500 result := targethasher.Result{ TargetNames: make([]string, numTargets), @@ -198,10 +198,10 @@ func TestResultToGetTargetGraphResponse_Chunking(t *testing.T) { result.Targets[name] = &targethasher.Target{Name: name, Hash: []byte{0}, RuleType: "go_library"} } - responses, err := ResultToGetTargetGraphResponse(context.Background(), result) + responses, err := ResultToGetTargetGraphResponse(context.Background(), result, 250, 50_000) require.NoError(t, err) - // 2 target chunks + 1 metadata chunk (500 targets well under DefaultMetadataMapChunkSize) + // 2 target chunks + 1 metadata chunk (500 targets well under metadataMapChunkSize) require.Len(t, responses, 3) for _, resp := range responses[:2] { diff --git a/example/cmd/query-bench/main.go b/example/cmd/query-bench/main.go index 3071b01..8efaa25 100644 --- a/example/cmd/query-bench/main.go +++ b/example/cmd/query-bench/main.go @@ -112,7 +112,7 @@ func run() error { totalDuration += elapsed fmt.Printf("run %d: targethasher: %v (%d targets)\n", i+1, elapsed.Round(time.Millisecond), len(targethasherResult.TargetNames)) start = time.Now() - response, err := common.ResultToGetTargetGraphResponse(ctx, targethasherResult) // also print the time to convert to GetTargetGraphResponse format + response, err := common.ResultToGetTargetGraphResponse(ctx, targethasherResult, 250, 50_000) if err != nil { return fmt.Errorf("run %d: converting to GetTargetGraphResponse: %w", i+1, err) } diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index f46c2ab..57ac3c1 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -206,7 +206,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget if err != nil { return nil, fmt.Errorf("compute target graph: %w", err) } - responses, err := common.ResultToGetTargetGraphResponse(ctx, result) + responses, err := common.ResultToGetTargetGraphResponse(ctx, result, b.config.Service.Chunking.TargetChunkSize, b.config.Service.Chunking.MetadataMapChunkSize) if err != nil { return nil, fmt.Errorf("convert target graph to response: %w", err) } From 799b96e902112c0af1c9674eae744df3b48ab830 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 17:31:59 -0700 Subject: [PATCH 4/7] Make sure config fields are all set in constructor --- config/config.go | 4 ++-- config/repository_config.go | 2 +- example/tango-config.yaml | 2 +- orchestrator/native_orchestrator.go | 2 +- orchestrator/testdata/config.yaml | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/config/config.go b/config/config.go index 9facddb..41decb7 100644 --- a/config/config.go +++ b/config/config.go @@ -85,8 +85,8 @@ func Parse(configFilePath string) (*Config, error) { if remote == "" { return nil, fmt.Errorf("repository[%d].remote must not be empty", i) } - if config.Repository[i].QueryTimeout <= 0 { - return nil, fmt.Errorf("repository[%d].query_timeout must be > 0", i) + if config.Repository[i].QueryTimeoutSeconds <= 0 { + return nil, fmt.Errorf("repository[%d].query_timeout_seconds must be > 0", i) } if _, exists := config.repositoryByRemote[remote]; exists { return nil, fmt.Errorf("duplicate repository remote %q", remote) diff --git a/config/repository_config.go b/config/repository_config.go index 4052ec4..8e23f0a 100644 --- a/config/repository_config.go +++ b/config/repository_config.go @@ -22,7 +22,7 @@ type RepositoryConfig struct { ExcludeExternalTargets bool `yaml:"exclude_external_targets"` BzlmodEnabled bool `yaml:"bzlmod_enabled"` BazelCommand string `yaml:"bazel_command"` - QueryTimeout int64 `yaml:"query_timeout"` // in seconds + QueryTimeoutSeconds int64 `yaml:"query_timeout_seconds"` BazelExtraArgs []string `yaml:"bazel_extra_args"` StreamBazelLogs bool `yaml:"stream_bazel_logs"` } diff --git a/example/tango-config.yaml b/example/tango-config.yaml index 1e45c0d..8608c65 100644 --- a/example/tango-config.yaml +++ b/example/tango-config.yaml @@ -15,7 +15,7 @@ repository: - "^@@?bazel_tools/" exclude_external_targets: true bzlmod_enabled: true - query_timeout: 300 + query_timeout_seconds: 300 # Service configuration service: diff --git a/orchestrator/native_orchestrator.go b/orchestrator/native_orchestrator.go index 57ac3c1..0f99272 100644 --- a/orchestrator/native_orchestrator.go +++ b/orchestrator/native_orchestrator.go @@ -187,7 +187,7 @@ func (b *nativeOrchestrator) GetTargetGraph(ctx context.Context, param GetTarget WorkspacePath: ws.Path(), Logger: b.logger, BazelCommand: repoCfg.BazelCommand, - QueryTimeout: time.Duration(repoCfg.QueryTimeout) * time.Second, + QueryTimeout: time.Duration(repoCfg.QueryTimeoutSeconds) * time.Second, StreamLogs: repoCfg.StreamBazelLogs, }) if err != nil { diff --git a/orchestrator/testdata/config.yaml b/orchestrator/testdata/config.yaml index 34887e0..2fc7793 100644 --- a/orchestrator/testdata/config.yaml +++ b/orchestrator/testdata/config.yaml @@ -7,7 +7,7 @@ repository: - "*.gen.go" exclude_external_targets: true bzlmod_enabled: true - query_timeout: 15 + query_timeout_seconds: 15 service: worker_pool_size: 3 From b5aeafc9376095c47741486126e344129fea6b32 Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 18:35:19 -0700 Subject: [PATCH 5/7] Update with tests --- config/BUILD.bazel | 9 +- config/config_test.go | 186 ++++++++++++++++++++++++++++++++++ controller/controller_test.go | 51 ++++++++++ core/bazel/bazel_test.go | 42 ++++++++ orchestrator/BUILD.bazel | 1 + 5 files changed, 288 insertions(+), 1 deletion(-) create mode 100644 config/config_test.go diff --git a/config/BUILD.bazel b/config/BUILD.bazel index 0c95044..84c243f 100644 --- a/config/BUILD.bazel +++ b/config/BUILD.bazel @@ -1,4 +1,4 @@ -load("@rules_go//go:def.bzl", "go_library") +load("@rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "config", @@ -12,3 +12,10 @@ go_library( visibility = ["//visibility:public"], deps = ["@com_github_goccy_go_yaml//:go-yaml"], ) + +go_test( + name = "config_test", + srcs = ["config_test.go"], + embed = [":config"], + deps = ["@com_github_stretchr_testify//require"], +) diff --git a/config/config_test.go b/config/config_test.go new file mode 100644 index 0000000..eff9be6 --- /dev/null +++ b/config/config_test.go @@ -0,0 +1,186 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// 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 config + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParse_Validation(t *testing.T) { + t.Parallel() + + validConfig := ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +storage: + type: memory +service: + worker_pool_size: 5 + repo_manager_clone_path: /tmp/tango + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +` + + tests := []struct { + name string + yaml string + wantErr bool + }{ + { + name: "valid config", + yaml: validConfig, + }, + { + name: "worker_pool_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "target_chunk_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "changed_target_chunk_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "metadata_map_chunk_size not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 +`, + wantErr: true, + }, + { + name: "repository remote not set", + yaml: ` +repository: + - query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "query_timeout_seconds not set", + yaml: ` +repository: + - remote: "git@github.com:org/repo" +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "duplicate repository remote", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + { + name: "worker_root_path without repo_manager_clone_path", + yaml: ` +repository: + - remote: "git@github.com:org/repo" + query_timeout_seconds: 300 +service: + worker_pool_size: 5 + worker_root_path: /tmp/workers + chunking: + target_chunk_size: 250 + changed_target_chunk_size: 125 + metadata_map_chunk_size: 50000 +`, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "config.yaml") + require.NoError(t, os.WriteFile(path, []byte(tt.yaml), 0o644)) + + cfg, err := Parse(path) + if tt.wantErr { + require.Error(t, err, "expected error when %s", tt.name) + } else { + require.NoError(t, err) + require.NotNil(t, cfg) + } + }) + } +} diff --git a/controller/controller_test.go b/controller/controller_test.go index 870b5c6..c591866 100644 --- a/controller/controller_test.go +++ b/controller/controller_test.go @@ -117,3 +117,54 @@ func TestLinkRequestCtx_CancelReleasesAfterFunc(t *testing.T) { cancelApp() assert.Equal(t, firstErr, linked.Err(), "linked.Err() must not change after cancelLink") } + +func TestNewController_Validation(t *testing.T) { + t.Parallel() + + valid := config.ChunkConfig{ + TargetChunkSize: 250, + ChangedTargetChunkSize: 125, + MetadataMapChunkSize: 50_000, + } + + tests := []struct { + name string + chunk config.ChunkConfig + wantErr bool + }{ + { + name: "valid", + chunk: valid, + }, + { + name: "target_chunk_size not set", + chunk: config.ChunkConfig{ChangedTargetChunkSize: 125, MetadataMapChunkSize: 50_000}, + wantErr: true, + }, + { + name: "changed_target_chunk_size not set", + chunk: config.ChunkConfig{TargetChunkSize: 250, MetadataMapChunkSize: 50_000}, + wantErr: true, + }, + { + name: "metadata_map_chunk_size not set", + chunk: config.ChunkConfig{TargetChunkSize: 250, ChangedTargetChunkSize: 125}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := NewController(context.Background(), Params{ + Logger: zap.NewNop(), + ChunkConfig: tt.chunk, + }) + if tt.wantErr { + require.Error(t, err, "expected error when %s", tt.name) + } else { + require.NoError(t, err) + } + }) + } +} diff --git a/core/bazel/bazel_test.go b/core/bazel/bazel_test.go index 409e918..c5cf88b 100644 --- a/core/bazel/bazel_test.go +++ b/core/bazel/bazel_test.go @@ -88,3 +88,45 @@ func TestNewBazelClient_WithNilExecCommand(t *testing.T) { assert.Equal(t, "/workspace", execCmd.Dir) assert.Contains(t, execCmd.Env, "KEY=value") } + +func TestNewBazelClient_Validation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + timeout time.Duration + wantErr bool + }{ + { + name: "valid timeout", + timeout: 5 * time.Minute, + }, + { + name: "query timeout not set", + timeout: 0, + wantErr: true, + }, + { + name: "negative query timeout", + timeout: -1 * time.Second, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := NewBazelClient(context.Background(), Params{ + BazelCommand: "bazel", + WorkspacePath: "/tmp/test", + Logger: zap.NewNop().Sugar(), + QueryTimeout: tt.timeout, + }) + if tt.wantErr { + require.Error(t, err, "expected error when %s", tt.name) + } else { + require.NoError(t, err) + } + }) + } +} diff --git a/orchestrator/BUILD.bazel b/orchestrator/BUILD.bazel index ed80f1d..0141613 100644 --- a/orchestrator/BUILD.bazel +++ b/orchestrator/BUILD.bazel @@ -44,6 +44,7 @@ go_test( "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", "@org_uber_go_mock//gomock", + "@org_uber_go_zap//:zap", "@org_uber_go_zap//zaptest", ], ) From 9316ed5fdadf4e32e8c24884b95549e2487c093e Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 19:13:55 -0700 Subject: [PATCH 6/7] Update --- core/bazel/query_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/core/bazel/query_test.go b/core/bazel/query_test.go index 2b72982..09337a8 100644 --- a/core/bazel/query_test.go +++ b/core/bazel/query_test.go @@ -57,7 +57,7 @@ func TestExecuteQuery_Success(t *testing.T) { mockCmd.EXPECT().Wait().Return(nil), ) client, err := NewBazelClient(context.Background(), Params{ - QueryTimeout: 5 * time.Minute, + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -102,7 +102,7 @@ func TestExecuteQuery_WithStartupOptions(t *testing.T) { mockCmd.EXPECT().Wait().Return(nil), ) client, err := NewBazelClient(context.Background(), Params{ - QueryTimeout: 5 * time.Minute, + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -232,7 +232,7 @@ func TestExecuteQueryInternal_Failures(t *testing.T) { tt.setupMock(mockCmd) client, err := NewBazelClient(context.Background(), Params{ - QueryTimeout: 5 * time.Minute, + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, @@ -262,7 +262,7 @@ func TestExecuteQuery_ErrorCase(t *testing.T) { mockCmd.EXPECT().StdoutPipe().Return(nil, errors.New("stdout pipe failed")) client, err := NewBazelClient(context.Background(), Params{ - QueryTimeout: 5 * time.Minute, + QueryTimeout: 5 * time.Minute, BazelCommand: "bazel", WorkspacePath: "/tmp/test", EnvVarsMap: map[string]string{}, From 29d855568f48558f44deb6c03678621378c4e30b Mon Sep 17 00:00:00 2001 From: Yushan Lin Date: Thu, 9 Jul 2026 19:14:18 -0700 Subject: [PATCH 7/7] Update --- orchestrator/BUILD.bazel | 1 - 1 file changed, 1 deletion(-) diff --git a/orchestrator/BUILD.bazel b/orchestrator/BUILD.bazel index 0141613..ed80f1d 100644 --- a/orchestrator/BUILD.bazel +++ b/orchestrator/BUILD.bazel @@ -44,7 +44,6 @@ go_test( "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", "@org_uber_go_mock//gomock", - "@org_uber_go_zap//:zap", "@org_uber_go_zap//zaptest", ], )