go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/analysis/internal/testutil/spantest.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 testutil
    16  
    17  import (
    18  	"context"
    19  	"os"
    20  	"path/filepath"
    21  	"testing"
    22  	"time"
    23  
    24  	"cloud.google.com/go/spanner"
    25  
    26  	"go.chromium.org/luci/common/errors"
    27  	"go.chromium.org/luci/common/spantest"
    28  	"go.chromium.org/luci/common/tsmon"
    29  	"go.chromium.org/luci/server/span"
    30  
    31  	. "github.com/smartystreets/goconvey/convey"
    32  )
    33  
    34  // cleanupDatabase deletes all data from all tables.
    35  func cleanupDatabase(ctx context.Context, client *spanner.Client) error {
    36  	_, err := client.Apply(ctx, []*spanner.Mutation{
    37  		// No need to explicitly delete interleaved tables.
    38  		spanner.Delete("ClusteringState", spanner.AllKeys()),
    39  		spanner.Delete("FailureAssociationRules", spanner.AllKeys()),
    40  		spanner.Delete("GerritChangelists", spanner.AllKeys()),
    41  		spanner.Delete("Ingestions", spanner.AllKeys()),
    42  		spanner.Delete("ReclusteringRuns", spanner.AllKeys()),
    43  		spanner.Delete("ReclusteringShards", spanner.AllKeys()),
    44  		spanner.Delete("TestResults", spanner.AllKeys()),
    45  		spanner.Delete("TestVariantRealms", spanner.AllKeys()),
    46  		spanner.Delete("TestRealms", spanner.AllKeys()),
    47  		spanner.Delete("TestVariantBranch", spanner.AllKeys()),
    48  		spanner.Delete("TestVariantBranchCheckpoint", spanner.AllKeys()),
    49  		spanner.Delete("Invocations", spanner.AllKeys()),
    50  	})
    51  	return err
    52  }
    53  
    54  // IntegrationTestContext returns a context for testing code that talks to Spanner
    55  // and uses tsmon.
    56  // Skips the test if integration tests are not enabled.
    57  //
    58  // Tests that use Spanner must not call t.Parallel().
    59  func IntegrationTestContext(tb testing.TB) context.Context {
    60  	// tsmon metrics are used fairly extensively throughout LUCI Analysis,
    61  	// especially in contexts that also use Spanner.
    62  	ctx := SpannerTestContext(tb)
    63  	ctx, _ = tsmon.WithDummyInMemory(ctx)
    64  	return ctx
    65  }
    66  
    67  // SpannerTestContext returns a context for testing code that talks to Spanner.
    68  // Skips the test if integration tests are not enabled.
    69  //
    70  // Tests that use Spanner must not call t.Parallel().
    71  func SpannerTestContext(tb testing.TB) context.Context {
    72  	return spantest.SpannerTestContext(tb, cleanupDatabase)
    73  }
    74  
    75  // findInitScript returns path //analysis/internal/span/init_db.sql.
    76  func findInitScript() (string, error) {
    77  	ancestor, err := filepath.Abs(".")
    78  	if err != nil {
    79  		return "", err
    80  	}
    81  
    82  	for {
    83  		scriptPath := filepath.Join(ancestor, "internal", "span", "init_db.sql")
    84  		_, err := os.Stat(scriptPath)
    85  		if os.IsNotExist(err) {
    86  			parent := filepath.Dir(ancestor)
    87  			if parent == ancestor {
    88  				return "", errors.Reason("init_db.sql not found").Err()
    89  			}
    90  			ancestor = parent
    91  			continue
    92  		}
    93  
    94  		return scriptPath, err
    95  	}
    96  }
    97  
    98  // SpannerTestMain is a test main function for packages that have tests that
    99  // talk to spanner. It creates/destroys a temporary spanner database
   100  // before/after running tests.
   101  //
   102  // This function never returns. Instead it calls os.Exit with the value returned
   103  // by m.Run().
   104  func SpannerTestMain(m *testing.M) {
   105  	spantest.SpannerTestMain(m, findInitScript)
   106  }
   107  
   108  // MustApply applies the mutations to the spanner client in the context.
   109  // Asserts that application succeeds.
   110  // Returns the commit timestamp.
   111  func MustApply(ctx context.Context, ms ...*spanner.Mutation) time.Time {
   112  	ct, err := span.Apply(ctx, ms)
   113  	So(err, ShouldBeNil)
   114  	return ct
   115  }