From 2244819507b06064bbafae1446e84080f9a3c8c2 Mon Sep 17 00:00:00 2001
From: Jackson Owens <jackson@cockroachlabs.com>
Date: Wed, 13 Nov 2024 14:10:35 -0500
Subject: [PATCH] storage: extend TestKeySchema_RandomKeys assertions

Assert that seeking to each key finds the right key.

Epic: none
Release note: none
---
 pkg/storage/pebble_key_schema_test.go | 38 +++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/pkg/storage/pebble_key_schema_test.go b/pkg/storage/pebble_key_schema_test.go
index 32b1af7e585f..b842a3795c26 100644
--- a/pkg/storage/pebble_key_schema_test.go
+++ b/pkg/storage/pebble_key_schema_test.go
@@ -340,6 +340,8 @@ func TestKeySchema_RandomKeys(t *testing.T) {
 	var it colblk.DataBlockIter
 	it.InitOnce(&keySchema, EngineKeyCompare, EngineKeySplit, nil)
 	require.NoError(t, it.Init(&dec, block.NoTransforms))
+	// Ensure that a scan across the block finds all the relevant keys.
+	var valBuf []byte
 	for k, kv := 0, it.First(); kv != nil; k, kv = k+1, it.Next() {
 		require.True(t, EngineKeyEqual(keys[k], kv.K.UserKey))
 		require.Zero(t, EngineKeyCompare(keys[k], kv.K.UserKey))
@@ -352,6 +354,42 @@ func TestKeySchema_RandomKeys(t *testing.T) {
 			t.Fatalf("key %q is longer than original key %q", kv.K.UserKey, keys[k])
 		}
 		checkEngineKey(kv.K.UserKey)
+
+		// We write keys[k] as the value too, so check that it's verbatim equal.
+		value, callerOwned, err := kv.V.Value(valBuf)
+		require.NoError(t, err)
+		require.Equal(t, keys[k], value)
+		if callerOwned {
+			valBuf = value
+		}
+	}
+	// Ensure that seeking to each key finds the key.
+	for i := range keys {
+		kv := it.SeekGE(keys[i], 0)
+		require.True(t, EngineKeyEqual(keys[i], kv.K.UserKey))
+		require.Zero(t, EngineKeyCompare(keys[i], kv.K.UserKey))
+	}
+	// Ensure seeking to just the prefix of each key finds a key with the same
+	// prefix.
+	for i := range keys {
+		si := EngineKeySplit(keys[i])
+		kv := it.SeekGE(keys[i][:si], 0)
+		require.True(t, EngineKeyEqual(keys[i][:si], pebble.Split(EngineKeySplit).Prefix(kv.K.UserKey)))
 	}
+	// Ensure seeking to the key but in random order finds the key.
+	for _, i := range rng.Perm(len(keys)) {
+		kv := it.SeekGE(keys[i], 0)
+		require.True(t, EngineKeyEqual(keys[i], kv.K.UserKey))
+		require.Zero(t, EngineKeyCompare(keys[i], kv.K.UserKey))
+
+		// We write keys[k] as the value too, so check that it's verbatim equal.
+		value, callerOwned, err := kv.V.Value(valBuf)
+		require.NoError(t, err)
+		require.Equal(t, keys[i], value)
+		if callerOwned {
+			valBuf = value
+		}
+	}
+
 	require.NoError(t, it.Close())
 }