kythe.io@v0.0.68-0.20240422202219-7225dbc01741/kythe/go/serving/pipeline/beamio/leveldb_test.go (about)

     1  /*
     2   * Copyright 2018 The Kythe Authors. All rights reserved.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *   http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package beamio
    18  
    19  import (
    20  	"io/ioutil"
    21  	"math/rand"
    22  	"os"
    23  	"sort"
    24  	"testing"
    25  	"unicode"
    26  
    27  	"github.com/apache/beam/sdks/go/pkg/beam"
    28  	"github.com/apache/beam/sdks/go/pkg/beam/testing/ptest"
    29  	"github.com/apache/beam/sdks/go/pkg/beam/transforms/stats"
    30  	"github.com/jmhodges/levigo"
    31  
    32  	_ "github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/local"
    33  )
    34  
    35  func TestLevelDBSink(t *testing.T) {
    36  	var strs []string
    37  	for i := 0; i < 28; i++ {
    38  		buf := ""
    39  		for c := 'A'; c <= 'z'; c++ {
    40  			if !unicode.IsLetter(c) {
    41  				continue
    42  			}
    43  			buf += string(c)
    44  			strs = append(strs, buf)
    45  		}
    46  		for len(buf) > 1 {
    47  			buf = buf[:len(buf)-1]
    48  			strs = append(strs, buf)
    49  		}
    50  	}
    51  
    52  	p, s, coll := ptest.CreateList(strs)
    53  	tbl := beam.ParDo(s, extendToKey, coll)
    54  
    55  	path, err := ioutil.TempDir("", "leveldb")
    56  	if err != nil {
    57  		t.Fatalf("Error creating temporary directory: %v", err)
    58  	}
    59  	defer os.RemoveAll(path)
    60  	WriteLevelDB(s, path, stats.Opts{NumQuantiles: 4, K: 100}, tbl)
    61  
    62  	if err := ptest.Run(p); err != nil {
    63  		t.Fatal(err)
    64  	}
    65  
    66  	opts := levigo.NewOptions()
    67  	defer opts.Close()
    68  	db, err := levigo.Open(path, opts)
    69  	if err != nil {
    70  		t.Fatal(err)
    71  	}
    72  	defer db.Close()
    73  
    74  	ropts := levigo.NewReadOptions()
    75  	defer ropts.Close()
    76  	iter := db.NewIterator(ropts)
    77  	defer iter.Close()
    78  
    79  	// "Randomly" seek for the known key-value entries
    80  	for seed := int64(0); seed < 10; seed++ {
    81  		rand.Seed(seed)
    82  		for i := 0; i < len(strs)*10; i++ {
    83  			s := strs[rand.Int()%len(strs)]
    84  			iter.Seek([]byte(s))
    85  			if !iter.Valid() {
    86  				t.Fatalf("Invalid Iterator: %v", iter.GetError())
    87  			} else if found := string(iter.Key()); found != s {
    88  				t.Errorf("Expected key %q; found %q", s, found)
    89  			} else if found := string(iter.Value()); found != s {
    90  				t.Errorf("Expected value %q; found %q", s, found)
    91  			}
    92  		}
    93  	}
    94  
    95  	// Seek for the known key-value entries in order.
    96  	sort.Strings(strs)
    97  	for _, s := range strs {
    98  		iter.Seek([]byte(s))
    99  		if !iter.Valid() {
   100  			t.Fatal("Invalid Iterator")
   101  		} else if found := string(iter.Key()); found != s {
   102  			t.Errorf("Expected key %q; found %q", s, found)
   103  		} else if found := string(iter.Value()); found != s {
   104  			t.Errorf("Expected value %q; found %q", s, found)
   105  		}
   106  	}
   107  }
   108  
   109  func extendToKey(v beam.T) (beam.T, beam.T) { return v, v }
   110  
   111  func TestSchemaPreservingPathJoin(t *testing.T) {
   112  	exp := "gs://foo/bar/baz"
   113  	res := schemePreservingPathJoin("gs://foo", "bar/baz")
   114  	if exp != res {
   115  		t.Fatalf("Expected [%s], got [%s]", exp, res)
   116  	}
   117  
   118  	exp = "foo/bar/baz"
   119  	res = schemePreservingPathJoin("foo//bar", "baz")
   120  	if exp != res {
   121  		t.Fatalf("Expected [%s], got [%s]", exp, res)
   122  	}
   123  
   124  	exp = "/foo/bar"
   125  	res = schemePreservingPathJoin("/foo/", "bar")
   126  	if exp != res {
   127  		t.Fatalf("Expected [%s], got [%s]", exp, res)
   128  	}
   129  }