github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/geo/geoindex/s2_geometry_index_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package geoindex
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/geo"
    19  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    20  	"github.com/cockroachdb/datadriven"
    21  )
    22  
    23  func TestS2GeometryIndexBasic(t *testing.T) {
    24  	defer leaktest.AfterTest(t)()
    25  
    26  	ctx := context.Background()
    27  	var index GeometryIndex
    28  	shapes := make(map[string]*geo.Geometry)
    29  	datadriven.RunTest(t, "testdata/s2_geometry", func(t *testing.T, d *datadriven.TestData) string {
    30  		switch d.Cmd {
    31  		case "init":
    32  			cfg := s2Config(t, d)
    33  			var minX, minY, maxX, maxY int
    34  			d.ScanArgs(t, "minx", &minX)
    35  			d.ScanArgs(t, "miny", &minY)
    36  			d.ScanArgs(t, "maxx", &maxX)
    37  			d.ScanArgs(t, "maxy", &maxY)
    38  			index = NewS2GeometryIndex(S2GeometryConfig{
    39  				MinX:     float64(minX),
    40  				MinY:     float64(minY),
    41  				MaxX:     float64(maxX),
    42  				MaxY:     float64(maxY),
    43  				S2Config: &cfg,
    44  			})
    45  			return ""
    46  		case "geometry":
    47  			g, err := geo.ParseGeometry(d.Input)
    48  			if err != nil {
    49  				return err.Error()
    50  			}
    51  			shapes[nameArg(t, d)] = g
    52  			return ""
    53  		case "index-keys":
    54  			return keysToString(index.InvertedIndexKeys(ctx, shapes[nameArg(t, d)]))
    55  		case "inner-covering":
    56  			return cellUnionToString(index.TestingInnerCovering(shapes[nameArg(t, d)]))
    57  		case "covers":
    58  			return spansToString(index.Covers(ctx, shapes[nameArg(t, d)]))
    59  		case "intersects":
    60  			return spansToString(index.Intersects(ctx, shapes[nameArg(t, d)]))
    61  		case "covered-by":
    62  			return checkExprAndToString(index.CoveredBy(ctx, shapes[nameArg(t, d)]))
    63  		default:
    64  			return fmt.Sprintf("unknown command: %s", d.Cmd)
    65  		}
    66  	})
    67  }