go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/analysis/internal/clustering/runs/testutils.go (about)

     1  // Copyright 2022 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package runs
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"cloud.google.com/go/spanner"
    22  
    23  	"go.chromium.org/luci/server/span"
    24  
    25  	"go.chromium.org/luci/analysis/internal/testutil"
    26  )
    27  
    28  const testProject = "myproject"
    29  
    30  // RunBuilder provides methods to build a reclustering run
    31  // for testing.
    32  type RunBuilder struct {
    33  	run         *ReclusteringRun
    34  	progressSet bool
    35  }
    36  
    37  // NewRun starts building a new Run, for testing.
    38  func NewRun(uniqifier int) *RunBuilder {
    39  	run := &ReclusteringRun{
    40  		Project:           testProject,
    41  		AttemptTimestamp:  time.Date(2010, time.January, 1, 1, 0, 0, uniqifier, time.UTC),
    42  		AlgorithmsVersion: int64(uniqifier + 1),
    43  		ConfigVersion:     time.Date(2011, time.January, 1, 1, 0, 0, uniqifier, time.UTC),
    44  		RulesVersion:      time.Date(2012, time.January, 1, 1, 0, 0, uniqifier, time.UTC),
    45  		ShardCount:        int64(uniqifier + 1),
    46  		ShardsReported:    int64(uniqifier / 2),
    47  		Progress:          int64(uniqifier) * 500,
    48  	}
    49  	return &RunBuilder{
    50  		run: run,
    51  	}
    52  }
    53  
    54  // WithProject specifies the project to use on the run.
    55  func (b *RunBuilder) WithProject(project string) *RunBuilder {
    56  	b.run.Project = project
    57  	return b
    58  }
    59  
    60  // WithAttemptTimestamp specifies the attempt timestamp to use on the run.
    61  func (b *RunBuilder) WithAttemptTimestamp(attemptTimestamp time.Time) *RunBuilder {
    62  	b.run.AttemptTimestamp = attemptTimestamp
    63  	return b
    64  }
    65  
    66  // WithRulesVersion specifies the rules version to use on the run.
    67  func (b *RunBuilder) WithRulesVersion(value time.Time) *RunBuilder {
    68  	b.run.RulesVersion = value
    69  	return b
    70  }
    71  
    72  // WithRulesVersion specifies the config version to use on the run.
    73  func (b *RunBuilder) WithConfigVersion(value time.Time) *RunBuilder {
    74  	b.run.ConfigVersion = value
    75  	return b
    76  }
    77  
    78  // WithAlgorithmsVersion specifies the algorithms version to use on the run.
    79  func (b *RunBuilder) WithAlgorithmsVersion(value int64) *RunBuilder {
    80  	b.run.AlgorithmsVersion = value
    81  	return b
    82  }
    83  
    84  // WithShardCount specifies the number of shards to use on the run.
    85  func (b *RunBuilder) WithShardCount(count int64) *RunBuilder {
    86  	if b.progressSet {
    87  		panic("Call WithShardCount before setting progress")
    88  	}
    89  	b.run.ShardCount = count
    90  	b.run.Progress = count * 500
    91  	return b
    92  }
    93  
    94  // WithNoProgress sets that no shards reported and no progress has been made.
    95  func (b *RunBuilder) WithNoReportedProgress() *RunBuilder {
    96  	b.run.ShardsReported = 0
    97  	b.run.Progress = 0
    98  	b.progressSet = true
    99  	return b
   100  }
   101  
   102  // WithReportedProgress sets that all shards have reported, and some progress
   103  // has been made. (progress is a value out of 1000 that is scaled to the
   104  // number of shards).
   105  func (b *RunBuilder) WithReportedProgress(progress int) *RunBuilder {
   106  	b.run.ShardsReported = b.run.ShardCount
   107  	b.run.Progress = b.run.ShardCount * int64(progress)
   108  	b.progressSet = true
   109  	return b
   110  }
   111  
   112  // WithCompletedProgress sets that all shards have reported, and the run
   113  // has completed.
   114  func (b *RunBuilder) WithCompletedProgress() *RunBuilder {
   115  	b.run.ShardsReported = b.run.ShardCount
   116  	b.run.Progress = b.run.ShardCount * 1000
   117  	b.progressSet = true
   118  	return b
   119  }
   120  
   121  func (b *RunBuilder) Build() *ReclusteringRun {
   122  	return b.run
   123  }
   124  
   125  // SetRunsForTesting replaces the set of stored runs to match the given set.
   126  func SetRunsForTesting(ctx context.Context, rs []*ReclusteringRun) error {
   127  	testutil.MustApply(ctx,
   128  		spanner.Delete("ReclusteringRuns", spanner.AllKeys()))
   129  	// Insert some ReclusteringRuns.
   130  	_, err := span.ReadWriteTransaction(ctx, func(ctx context.Context) error {
   131  		for _, r := range rs {
   132  			if err := Create(ctx, r); err != nil {
   133  				return err
   134  			}
   135  		}
   136  		return nil
   137  	})
   138  	return err
   139  }