go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/teams/internal/testutil/spantest.go (about)

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