package test import ( "bytes" "context" "fmt" "io" "strings" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/grafana/grafana/pkg/storage/unified/resource" "github.com/grafana/grafana/pkg/util/testutil" ) // Test names for the KV test suite const ( TestKVGet = "get operations" TestKVSave = "save operations" TestKVDelete = "delete operations" TestKVKeys = "keys listing" TestKVKeysWithLimits = "keys with limits and ranges" TestKVKeysWithSort = "keys with sorting" TestKVConcurrent = "concurrent operations" TestKVUnixTimestamp = "unix timestamp" TestKVBatchGet = "batch get operations" TestKVBatchDelete = "batch delete operations" ) // NewKVFunc is a function that creates a new KV instance for testing type NewKVFunc func(ctx context.Context) resource.KV // KVTestOptions configures which tests to run type KVTestOptions struct { NSPrefix string // namespace prefix for isolation } // GenerateRandomKVPrefix creates a random namespace prefix for test isolation func GenerateRandomKVPrefix() string { return fmt.Sprintf("kvtest-%d", time.Now().UnixNano()) } // RunKVTest runs the KV test suite func RunKVTest(t *testing.T, newKV NewKVFunc, opts *KVTestOptions) { if opts == nil { opts = &KVTestOptions{} } if opts.NSPrefix == "" { opts.NSPrefix = GenerateRandomKVPrefix() } t.Logf("Running KV tests with namespace prefix: %s", opts.NSPrefix) cases := []struct { name string fn func(*testing.T, resource.KV, string) }{ {TestKVGet, runTestKVGet}, {TestKVSave, runTestKVSave}, {TestKVDelete, runTestKVDelete}, {TestKVKeys, runTestKVKeys}, {TestKVKeysWithLimits, runTestKVKeysWithLimits}, {TestKVKeysWithSort, runTestKVKeysWithSort}, {TestKVConcurrent, runTestKVConcurrent}, {TestKVUnixTimestamp, runTestKVUnixTimestamp}, {TestKVBatchGet, runTestKVBatchGet}, {TestKVBatchDelete, runTestKVBatchDelete}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { tc.fn(t, newKV(context.Background()), opts.NSPrefix) }) } } func runTestKVGet(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-get" t.Run("get existing key", func(t *testing.T) { // First save a key testValue := "test value for get" saveKVHelper(t, kv, ctx, section, "existing-key", strings.NewReader(testValue)) // Now get it reader, err := kv.Get(ctx, section, "existing-key") require.NoError(t, err) // Read the value value, err := io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, testValue, string(value)) // Close the value reader err = reader.Close() require.NoError(t, err) }) t.Run("get non-existent key", func(t *testing.T) { _, err := kv.Get(ctx, section, "non-existent-key") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) }) t.Run("get with empty section", func(t *testing.T) { _, err := kv.Get(ctx, "", "some-key") assert.Error(t, err) assert.Contains(t, err.Error(), "section is required") }) } func runTestKVSave(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-save" t.Run("save new key", func(t *testing.T) { testValue := "new test value" saveKVHelper(t, kv, ctx, section, "new-key", strings.NewReader(testValue)) // Verify it was saved reader, err := kv.Get(ctx, section, "new-key") require.NoError(t, err) value, err := io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, testValue, string(value)) err = reader.Close() require.NoError(t, err) }) t.Run("save overwrite existing key", func(t *testing.T) { // First save saveKVHelper(t, kv, ctx, section, "overwrite-key", strings.NewReader("old value")) // Overwrite newValue := "new value" saveKVHelper(t, kv, ctx, section, "overwrite-key", strings.NewReader(newValue)) // Verify it was updated reader, err := kv.Get(ctx, section, "overwrite-key") require.NoError(t, err) value, err := io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, newValue, string(value)) err = reader.Close() require.NoError(t, err) }) t.Run("save with empty section", func(t *testing.T) { _, err := kv.Save(ctx, "", "some-key") assert.Error(t, err) assert.Contains(t, err.Error(), "section is required") }) t.Run("save binary data", func(t *testing.T) { binaryData := []byte{0x00, 0x01, 0x02, 0x03, 0xFF, 0xFE, 0xFD} saveKVHelper(t, kv, ctx, section, "binary-key", bytes.NewReader(binaryData)) // Verify binary data reader, err := kv.Get(ctx, section, "binary-key") require.NoError(t, err) value, err := io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, binaryData, value) err = reader.Close() require.NoError(t, err) }) t.Run("save key with no data", func(t *testing.T) { // Save a key with empty data saveKVHelper(t, kv, ctx, section, "empty-key", strings.NewReader("")) // Verify it was saved with empty data reader, err := kv.Get(ctx, section, "empty-key") require.NoError(t, err) value, err := io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, "", string(value)) assert.Len(t, value, 0) err = reader.Close() require.NoError(t, err) }) } func runTestKVDelete(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-delete" t.Run("delete existing key", func(t *testing.T) { // First create a key saveKVHelper(t, kv, ctx, section, "delete-key", strings.NewReader("delete me")) // Verify it exists _, err := kv.Get(ctx, section, "delete-key") require.NoError(t, err) // Delete it err = kv.Delete(ctx, section, "delete-key") require.NoError(t, err) // Verify it's gone _, err = kv.Get(ctx, section, "delete-key") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) }) t.Run("delete non-existent key", func(t *testing.T) { err := kv.Delete(ctx, section, "non-existent-delete-key") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) }) t.Run("delete with empty section", func(t *testing.T) { err := kv.Delete(ctx, "", "some-key") assert.Error(t, err) assert.Contains(t, err.Error(), "section is required") }) } func runTestKVKeys(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-keys" // Setup test data testKeys := []string{"a1", "a2", "b1", "b2", "c1"} for _, key := range testKeys { saveKVHelper(t, kv, ctx, section, key, strings.NewReader("value"+key)) } t.Run("list all keys", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{}) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, testKeys, keys) }) t.Run("list keys with empty section", func(t *testing.T) { var keys []string var errors []error for k, err := range kv.Keys(ctx, "", resource.ListOptions{}) { if err != nil { errors = append(errors, err) break } keys = append(keys, k) } assert.Len(t, errors, 1) assert.Contains(t, errors[0].Error(), "section is required") assert.Empty(t, keys) }) t.Run("list keys returns 0 keys", func(t *testing.T) { // Use a different section with no keys emptySection := nsPrefix + "-empty-keys" var keys []string for k, err := range kv.Keys(ctx, emptySection, resource.ListOptions{}) { require.NoError(t, err) keys = append(keys, k) } assert.Empty(t, keys) assert.Len(t, keys, 0) }) } func runTestKVKeysWithLimits(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-keys-limits" // Setup test data testKeys := []string{"a1", "a2", "b1", "b2", "c1", "c2", "d1", "d2"} for _, key := range testKeys { saveKVHelper(t, kv, ctx, section, key, strings.NewReader("value"+key)) } t.Run("keys with limit", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{Limit: 3}) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"a1", "a2", "b1"}, keys) }) t.Run("keys with range", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{StartKey: "b", EndKey: "d"}) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"b1", "b2", "c1", "c2"}, keys) }) t.Run("keys with prefix", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{ StartKey: "c", EndKey: resource.PrefixRangeEnd("c"), }) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"c1", "c2"}, keys) }) t.Run("keys with limit and range", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{ StartKey: "a", EndKey: "c", Limit: 2, }) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"a1", "a2"}, keys) }) } func runTestKVKeysWithSort(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-keys-sort" // Setup test data testKeys := []string{"a1", "a2", "b1", "b2", "c1"} for _, key := range testKeys { saveKVHelper(t, kv, ctx, section, key, strings.NewReader("value"+key)) } t.Run("keys in ascending order (default)", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{Sort: resource.SortOrderAsc}) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"a1", "a2", "b1", "b2", "c1"}, keys) }) t.Run("keys in descending order", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{Sort: resource.SortOrderDesc}) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"c1", "b2", "b1", "a2", "a1"}, keys) }) t.Run("keys descending with prefix", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{ StartKey: "a", EndKey: resource.PrefixRangeEnd("a"), Sort: resource.SortOrderDesc, }) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"a2", "a1"}, keys) }) t.Run("keys descending with limit", func(t *testing.T) { var keys []string for k, err := range kv.Keys(ctx, section, resource.ListOptions{ Sort: resource.SortOrderDesc, Limit: 3, }) { require.NoError(t, err) keys = append(keys, k) } assert.Equal(t, []string{"c1", "b2", "b1"}, keys) }) } func runTestKVConcurrent(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(60*time.Second)) section := nsPrefix + "-concurrent" t.Run("concurrent save and get operations", func(t *testing.T) { const numGoroutines = 10 const numOperations = 20 done := make(chan error, numGoroutines) for i := 0; i < numGoroutines; i++ { go func(goroutineID int) { var err error defer func() { done <- err }() for j := 0; j < numOperations; j++ { key := fmt.Sprintf("concurrent-key-%d-%d", goroutineID, j) value := fmt.Sprintf("concurrent-value-%d-%d", goroutineID, j) // Save writer, err := kv.Save(ctx, section, key) if err != nil { return } defer func() { err := writer.Close() require.NoError(t, err) }() _, err = io.Copy(writer, strings.NewReader(value)) if err != nil { return } err = writer.Close() if err != nil { return } // Get immediately reader, err := kv.Get(ctx, section, key) if err != nil { return } readValue, err := io.ReadAll(reader) require.NoError(t, err) err = reader.Close() require.NoError(t, err) assert.Equal(t, value, string(readValue)) } }(i) } // Wait for all goroutines to complete for i := 0; i < numGoroutines; i++ { err := <-done require.NoError(t, err) } }) t.Run("concurrent save, delete, and list operations", func(t *testing.T) { const numGoroutines = 5 done := make(chan error, numGoroutines) for i := 0; i < numGoroutines; i++ { go func(goroutineID int) { var err error defer func() { done <- err }() key := fmt.Sprintf("concurrent-ops-key-%d", goroutineID) value := fmt.Sprintf("concurrent-ops-value-%d", goroutineID) // Save writer, err := kv.Save(ctx, section, key) if err != nil { return } defer func() { err := writer.Close() require.NoError(t, err) }() _, err = io.Copy(writer, strings.NewReader(value)) if err != nil { return } err = writer.Close() if err != nil { return } // List to verify it exists found := false for k, err := range kv.Keys(ctx, section, resource.ListOptions{}) { if err != nil { return } if k == key { found = true break } } if !found { err = fmt.Errorf("key %s not found in list", key) return } // Delete err = kv.Delete(ctx, section, key) if err != nil { return } // Verify it's deleted _, err = kv.Get(ctx, section, key) require.ErrorIs(t, resource.ErrNotFound, err) err = nil // Expected error, so clear it }(i) } // Wait for all goroutines to complete for i := 0; i < numGoroutines; i++ { err := <-done require.NoError(t, err) } }) } func runTestKVUnixTimestamp(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) t.Run("unix timestamp returns reasonable value", func(t *testing.T) { timestamp, err := kv.UnixTimestamp(ctx) require.NoError(t, err) now := time.Now().Unix() // Allow for some time difference (up to 5 seconds) assert.InDelta(t, now, timestamp, 5) }) t.Run("unix timestamp is consistent", func(t *testing.T) { timestamp1, err := kv.UnixTimestamp(ctx) require.NoError(t, err) timestamp2, err := kv.UnixTimestamp(ctx) require.NoError(t, err) // Should be very close (within 1 second) require.InDelta(t, timestamp1, timestamp2, 1) }) } func runTestKVBatchGet(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-batchget" t.Run("batch get existing keys", func(t *testing.T) { // Setup test data testData := map[string]string{ "key1": "value1", "key2": "value2", "key3": "value3", } // Save test data for key, value := range testData { saveKVHelper(t, kv, ctx, section, key, strings.NewReader(value)) } // Batch get all keys keys := []string{"key1", "key2", "key3"} type result struct { key string value string } var results []result for kv, err := range kv.BatchGet(ctx, section, keys) { require.NoError(t, err) value, err := io.ReadAll(kv.Value) require.NoError(t, err) err = kv.Value.Close() require.NoError(t, err) results = append(results, result{key: kv.Key, value: string(value)}) } // Verify results assert.Len(t, results, 3) // Check that all keys are present and in order expectedKeys := []string{"key1", "key2", "key3"} actualKeys := make([]string, len(results)) for i, r := range results { actualKeys[i] = r.key } assert.Equal(t, expectedKeys, actualKeys) // Verify values for _, r := range results { assert.Equal(t, testData[r.key], r.value) } }) t.Run("batch get with non-existent keys", func(t *testing.T) { // Setup some test data saveKVHelper(t, kv, ctx, section, "existing-key", strings.NewReader("existing-value")) // Batch get with mix of existing and non-existent keys keys := []string{"existing-key", "non-existent-1", "non-existent-2"} type result struct { key string value string } var results []result for kv, err := range kv.BatchGet(ctx, section, keys) { require.NoError(t, err) value, err := io.ReadAll(kv.Value) require.NoError(t, err) err = kv.Value.Close() require.NoError(t, err) results = append(results, result{key: kv.Key, value: string(value)}) } // Should only return the existing key assert.Len(t, results, 1) assert.Equal(t, "existing-key", results[0].key) assert.Equal(t, "existing-value", results[0].value) }) t.Run("batch get with all non-existent keys", func(t *testing.T) { keys := []string{"non-existent-1", "non-existent-2", "non-existent-3"} var results []resource.KeyValue for kv, err := range kv.BatchGet(ctx, section, keys) { require.NoError(t, err) results = append(results, kv) } // Should return no results assert.Empty(t, results) }) t.Run("batch get with empty keys list", func(t *testing.T) { keys := []string{} var results []resource.KeyValue for kv, err := range kv.BatchGet(ctx, section, keys) { require.NoError(t, err) results = append(results, kv) } // Should return no results assert.Empty(t, results) }) t.Run("batch get with empty section", func(t *testing.T) { keys := []string{"some-key"} var errors []error for kv, err := range kv.BatchGet(ctx, "", keys) { if err != nil { errors = append(errors, err) break } _ = kv // unused } assert.Len(t, errors, 1) assert.Contains(t, errors[0].Error(), "section is required") }) t.Run("batch get preserves order", func(t *testing.T) { // Setup test data testData := map[string]string{ "z-key": "z-value", "a-key": "a-value", "m-key": "m-value", } // Save test data for key, value := range testData { saveKVHelper(t, kv, ctx, section, key, strings.NewReader(value)) } // Batch get in specific order keys := []string{"z-key", "a-key", "m-key"} var results []string for kv, err := range kv.BatchGet(ctx, section, keys) { require.NoError(t, err) err = kv.Value.Close() require.NoError(t, err) results = append(results, kv.Key) } // Verify order is preserved assert.Len(t, results, 3) expectedOrder := []string{"z-key", "a-key", "m-key"} assert.Equal(t, expectedOrder, results) }) } func runTestKVBatchDelete(t *testing.T, kv resource.KV, nsPrefix string) { ctx := testutil.NewTestContext(t, time.Now().Add(30*time.Second)) section := nsPrefix + "-batchdelete" t.Run("batch delete existing keys", func(t *testing.T) { // Setup test data testData := map[string]string{ "key1": "value1", "key2": "value2", "key3": "value3", } // Save test data for key, value := range testData { saveKVHelper(t, kv, ctx, section, key, strings.NewReader(value)) } // Verify keys exist before deletion for key := range testData { _, err := kv.Get(ctx, section, key) require.NoError(t, err) } // Batch delete all keys keys := []string{"key1", "key2", "key3"} err := kv.BatchDelete(ctx, section, keys) require.NoError(t, err) // Verify all keys are deleted for _, key := range keys { _, err := kv.Get(ctx, section, key) assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) } }) t.Run("batch delete with non-existent keys", func(t *testing.T) { // Setup some test data saveKVHelper(t, kv, ctx, section, "existing-key-1", strings.NewReader("value1")) saveKVHelper(t, kv, ctx, section, "existing-key-2", strings.NewReader("value2")) // Batch delete with mix of existing and non-existent keys keys := []string{"existing-key-1", "non-existent-1", "existing-key-2", "non-existent-2"} err := kv.BatchDelete(ctx, section, keys) require.NoError(t, err) // Verify existing keys are deleted _, err = kv.Get(ctx, section, "existing-key-1") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) _, err = kv.Get(ctx, section, "existing-key-2") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) }) t.Run("batch delete with all non-existent keys", func(t *testing.T) { // Batch delete keys that don't exist keys := []string{"non-existent-1", "non-existent-2", "non-existent-3"} err := kv.BatchDelete(ctx, section, keys) require.NoError(t, err) }) t.Run("batch delete with empty keys list", func(t *testing.T) { keys := []string{} err := kv.BatchDelete(ctx, section, keys) require.NoError(t, err) }) t.Run("batch delete with empty section", func(t *testing.T) { keys := []string{"some-key"} err := kv.BatchDelete(ctx, "", keys) assert.Error(t, err) assert.Contains(t, err.Error(), "section is required") }) t.Run("batch delete preserves other keys", func(t *testing.T) { // Setup test data saveKVHelper(t, kv, ctx, section, "keep-key-1", strings.NewReader("keep-value-1")) saveKVHelper(t, kv, ctx, section, "delete-key-1", strings.NewReader("delete-value-1")) saveKVHelper(t, kv, ctx, section, "keep-key-2", strings.NewReader("keep-value-2")) saveKVHelper(t, kv, ctx, section, "delete-key-2", strings.NewReader("delete-value-2")) // Batch delete specific keys keys := []string{"delete-key-1", "delete-key-2"} err := kv.BatchDelete(ctx, section, keys) require.NoError(t, err) // Verify deleted keys are gone _, err = kv.Get(ctx, section, "delete-key-1") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) _, err = kv.Get(ctx, section, "delete-key-2") assert.Error(t, err) assert.Equal(t, resource.ErrNotFound, err) // Verify kept keys still exist reader, err := kv.Get(ctx, section, "keep-key-1") require.NoError(t, err) value, err := io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, "keep-value-1", string(value)) err = reader.Close() require.NoError(t, err) reader, err = kv.Get(ctx, section, "keep-key-2") require.NoError(t, err) value, err = io.ReadAll(reader) require.NoError(t, err) assert.Equal(t, "keep-value-2", string(value)) err = reader.Close() require.NoError(t, err) }) } // saveKVHelper is a helper function to save data to KV store using the new WriteCloser interface func saveKVHelper(t *testing.T, kv resource.KV, ctx context.Context, section, key string, value io.Reader) { t.Helper() writer, err := kv.Save(ctx, section, key) require.NoError(t, err) _, err = io.Copy(writer, value) require.NoError(t, err) err = writer.Close() require.NoError(t, err) }