go.mercari.io/datastore@v1.8.2/internal/testutils/testutils.go (about)

     1  package testutils
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"testing"
     7  
     8  	"github.com/favclip/testerator/v2"
     9  	"go.mercari.io/datastore"
    10  	"go.mercari.io/datastore/aedatastore"
    11  	"go.mercari.io/datastore/clouddatastore"
    12  	"google.golang.org/api/iterator"
    13  )
    14  
    15  // EmitCleanUpLog is flag for emit Datastore clean up log.
    16  var EmitCleanUpLog = false
    17  
    18  // SetupCloudDatastore returns CloudDatastore clients and function for cleaning.
    19  func SetupCloudDatastore(t *testing.T) (context.Context, datastore.Client, func()) {
    20  	ctx := context.Background()
    21  	client, err := clouddatastore.FromContext(ctx)
    22  	if err != nil {
    23  		t.Fatal(err)
    24  	}
    25  
    26  	return ctx, client, func() {
    27  		defer client.Close()
    28  
    29  		q := client.NewQuery("__kind__").KeysOnly()
    30  		keys, err := client.GetAll(ctx, q, nil)
    31  		if err != nil {
    32  			t.Fatal(err)
    33  		}
    34  		if len(keys) == 0 {
    35  			return
    36  		}
    37  
    38  		kinds := make([]string, 0, len(keys))
    39  		for _, key := range keys {
    40  			kinds = append(kinds, key.Name())
    41  		}
    42  
    43  		if EmitCleanUpLog {
    44  			t.Logf("remove %s", strings.Join(kinds, ", "))
    45  		}
    46  
    47  		for _, kind := range kinds {
    48  
    49  			cnt := 0
    50  			for {
    51  				q := client.NewQuery(kind).Limit(1000).KeysOnly()
    52  				keys, err := client.GetAll(ctx, q, nil)
    53  				if err != nil {
    54  					t.Fatal(err)
    55  				}
    56  				err = client.DeleteMulti(ctx, keys)
    57  				if err != nil {
    58  					t.Fatal(err)
    59  				}
    60  
    61  				cnt += len(keys)
    62  
    63  				if len(keys) != 1000 {
    64  					if EmitCleanUpLog {
    65  						t.Logf("remove %s entity: %d", kind, cnt)
    66  					}
    67  					break
    68  				}
    69  			}
    70  		}
    71  	}
    72  }
    73  
    74  // SetupAEDatastore returns AEDatastore clients and function for cleaning.
    75  func SetupAEDatastore(t *testing.T) (context.Context, datastore.Client, func()) {
    76  	_, ctx, err := testerator.SpinUp()
    77  	if err != nil {
    78  		t.Fatal(err.Error())
    79  	}
    80  
    81  	client, err := aedatastore.FromContext(ctx)
    82  	if err != nil {
    83  		t.Fatal(err)
    84  	}
    85  
    86  	return ctx, client, func() { testerator.SpinDown() }
    87  }
    88  
    89  // CleanUpAllEntities in Datastore
    90  func CleanUpAllEntities(ctx context.Context, client datastore.Client) {
    91  	q := client.NewQuery("__kind__").KeysOnly()
    92  	iter := client.Run(ctx, q)
    93  	var kinds []string
    94  	for {
    95  		key, err := iter.Next(nil)
    96  		if err == iterator.Done {
    97  			break
    98  		}
    99  		if err != nil {
   100  			panic(err)
   101  		}
   102  		kinds = append(kinds, key.Name())
   103  	}
   104  
   105  	for _, kind := range kinds {
   106  		q := client.NewQuery(kind).KeysOnly()
   107  		keys, err := client.GetAll(ctx, q, nil)
   108  		if err != nil {
   109  			panic(err)
   110  		}
   111  		err = client.DeleteMulti(ctx, keys)
   112  		if err != nil {
   113  			panic(err)
   114  		}
   115  	}
   116  }