mirror of
https://github.com/gofiber/storage.git
synced 2025-10-04 16:22:52 +08:00
Refactor and optimize code
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -51,63 +52,72 @@ func TestCassandraStorage(t *testing.T) {
|
||||
// Test cases
|
||||
t.Run("KeyspaceCreation", func(t *testing.T) {
|
||||
connectionURL := newTestStore(t)
|
||||
store := New(Config{
|
||||
store, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_keyspace_creation",
|
||||
Table: "test_kv",
|
||||
Expiration: 5 * time.Second,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
testKeyspaceCreation(t, connectionURL, store)
|
||||
})
|
||||
|
||||
t.Run("BasicOperations", func(t *testing.T) {
|
||||
connectionURL := newTestStore(t)
|
||||
store := New(Config{
|
||||
store, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_basic_ops",
|
||||
Table: "test_kv",
|
||||
Expiration: 5 * time.Second,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
testBasicOperations(t, store)
|
||||
})
|
||||
|
||||
t.Run("ExpirableKeys", func(t *testing.T) {
|
||||
connectionURL := newTestStore(t)
|
||||
store := New(Config{
|
||||
store, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_expirable",
|
||||
Table: "test_kv",
|
||||
Expiration: 5 * time.Second,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
testExpirableKeys(t, store)
|
||||
})
|
||||
|
||||
t.Run("ConcurrentAccess", func(t *testing.T) {
|
||||
connectionURL := newTestStore(t)
|
||||
store := New(Config{
|
||||
store, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_concurrent",
|
||||
Table: "test_kv",
|
||||
Expiration: 5 * time.Second,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
testConcurrentAccess(t, store)
|
||||
})
|
||||
|
||||
t.Run("Reset", func(t *testing.T) {
|
||||
connectionURL := newTestStore(t)
|
||||
store := New(Config{
|
||||
store, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_reset",
|
||||
Table: "test_kv",
|
||||
Expiration: 5 * time.Second,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer store.Close()
|
||||
testReset(t, connectionURL, store)
|
||||
})
|
||||
|
||||
t.Run("IdentifierValidation", func(t *testing.T) {
|
||||
testIdentifierValidation(t)
|
||||
})
|
||||
}
|
||||
|
||||
// testKeyspaceCreation tests the keyspace creation functionality.
|
||||
@@ -171,7 +181,6 @@ func testBasicOperations(t *testing.T, store *Storage) {
|
||||
|
||||
// testExpirableKeys tests the expirable keys functionality.
|
||||
func testExpirableKeys(t *testing.T, store *Storage) {
|
||||
|
||||
// Set keys with different expiration settings
|
||||
// Key with default TTL (exp = 0 means use default)
|
||||
err := store.Set("key_default_ttl", []byte("value1"), 0)
|
||||
@@ -199,37 +208,37 @@ func testExpirableKeys(t *testing.T, store *Storage) {
|
||||
require.Equal(t, []byte("value3"), value)
|
||||
|
||||
// Wait for specific TTL to expire
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
|
||||
// Specific TTL key should be gone, others should remain
|
||||
value, err = store.Get("key_specific_ttl")
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, value, "Key with 1s TTL should have expired")
|
||||
require.Eventually(t, func() bool {
|
||||
v, _ := store.Get("key_specific_ttl")
|
||||
return v == nil
|
||||
}, 3*time.Second, 100*time.Millisecond,
|
||||
"Key with 1s TTL should have expired")
|
||||
|
||||
// Default TTL key should still exist
|
||||
value, err = store.Get("key_default_ttl")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("value1"), value, "Key with default TTL should still exist")
|
||||
|
||||
// No TTL key should still exist
|
||||
value, err = store.Get("key_no_ttl")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("value3"), value, "Key with no TTL should still exist")
|
||||
|
||||
// Wait for default TTL to expire
|
||||
time.Sleep(4 * time.Second)
|
||||
|
||||
// Default TTL key should be gone, no TTL key should remain
|
||||
value, err = store.Get("key_default_ttl")
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, value, "Key with default TTL should have expired")
|
||||
require.Eventually(t, func() bool {
|
||||
v, _ := store.Get("key_default_ttl")
|
||||
return v == nil
|
||||
}, 6*time.Second, 100*time.Millisecond,
|
||||
"Key with default TTL should have expired")
|
||||
|
||||
// No TTL key should still exist
|
||||
value, err = store.Get("key_no_ttl")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("value3"), value, "Key with no TTL should still exist")
|
||||
}
|
||||
|
||||
// / testReset tests the Reset method.
|
||||
// testReset tests the Reset method.
|
||||
func testReset(t *testing.T, connectionURL string, store *Storage) {
|
||||
|
||||
// Set some keys
|
||||
err := store.Set("key1", []byte("value1"), 0)
|
||||
require.NoError(t, err)
|
||||
@@ -255,35 +264,40 @@ func testReset(t *testing.T, connectionURL string, store *Storage) {
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, value, "Key should be deleted after reset")
|
||||
|
||||
// Close the first store before creating a new one
|
||||
store.Close()
|
||||
|
||||
// Create new storage with Reset flag
|
||||
store = New(Config{
|
||||
newStore, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_reset",
|
||||
Table: "test_kv",
|
||||
Reset: true,
|
||||
})
|
||||
defer store.Close()
|
||||
require.NoError(t, err)
|
||||
defer newStore.Close()
|
||||
|
||||
// Set a key
|
||||
err = store.Set("key3", []byte("value3"), 0)
|
||||
err = newStore.Set("key3", []byte("value3"), 0)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify key exists
|
||||
value, err = store.Get("key3")
|
||||
value, err = newStore.Get("key3")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte("value3"), value)
|
||||
}
|
||||
|
||||
// testConcurrentAccess tests concurrent access to the storage.
|
||||
func testConcurrentAccess(t *testing.T, store *Storage) {
|
||||
|
||||
// Number of goroutines
|
||||
const concurrentOps = 10
|
||||
done := make(chan bool, concurrentOps)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(concurrentOps)
|
||||
|
||||
// Run concurrent operations
|
||||
for i := 0; i < concurrentOps; i++ {
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
// Set key
|
||||
key := fmt.Sprintf("key%d", id)
|
||||
value := []byte(fmt.Sprintf("value%d", id))
|
||||
@@ -303,33 +317,28 @@ func testConcurrentAccess(t *testing.T, store *Storage) {
|
||||
retrievedValue, err = store.Get(key)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, retrievedValue)
|
||||
|
||||
done <- true
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Wait for all goroutines to complete
|
||||
for i := 0; i < concurrentOps; i++ {
|
||||
<-done
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func Benchmark_Cassandra_Set(b *testing.B) {
|
||||
|
||||
connectionURL := newTestStore(b)
|
||||
|
||||
// Create new storage
|
||||
store := New(Config{
|
||||
store, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_concurrent",
|
||||
Table: "test_kv",
|
||||
})
|
||||
require.NoError(b, err)
|
||||
defer store.Close()
|
||||
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
|
||||
var err error
|
||||
for i := 0; i < b.N; i++ {
|
||||
err = store.Set("john", []byte("doe"), 0)
|
||||
}
|
||||
@@ -338,21 +347,22 @@ func Benchmark_Cassandra_Set(b *testing.B) {
|
||||
}
|
||||
|
||||
func Benchmark_Cassandra_Get(b *testing.B) {
|
||||
|
||||
connectionURL := newTestStore(b)
|
||||
|
||||
// Create new storage
|
||||
client := New(Config{
|
||||
client, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_concurrent",
|
||||
Table: "test_kv",
|
||||
})
|
||||
require.NoError(b, err)
|
||||
defer client.Close()
|
||||
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
|
||||
err := client.Set("john", []byte("doe"), 0)
|
||||
err = client.Set("john", []byte("doe"), 0)
|
||||
require.NoError(b, err)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, err = client.Get("john")
|
||||
@@ -362,21 +372,20 @@ func Benchmark_Cassandra_Get(b *testing.B) {
|
||||
}
|
||||
|
||||
func Benchmark_Cassandra_Set_And_Delete(b *testing.B) {
|
||||
|
||||
connectionURL := newTestStore(b)
|
||||
|
||||
// Create new storage
|
||||
client := New(Config{
|
||||
client, err := New(Config{
|
||||
Hosts: []string{connectionURL},
|
||||
Keyspace: "test_concurrent",
|
||||
Table: "test_kv",
|
||||
})
|
||||
require.NoError(b, err)
|
||||
defer client.Close()
|
||||
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
|
||||
var err error
|
||||
for i := 0; i < b.N; i++ {
|
||||
_ = client.Set("john", []byte("doe"), 0)
|
||||
err = client.Delete("john")
|
||||
@@ -384,3 +393,47 @@ func Benchmark_Cassandra_Set_And_Delete(b *testing.B) {
|
||||
|
||||
require.NoError(b, err)
|
||||
}
|
||||
|
||||
// testIdentifierValidation tests the validateIdentifier function
|
||||
func testIdentifierValidation(t *testing.T) {
|
||||
// Test valid identifiers
|
||||
validCases := []string{
|
||||
"test",
|
||||
"test123",
|
||||
"test_123",
|
||||
"TEST",
|
||||
"Test123",
|
||||
"test_table",
|
||||
}
|
||||
|
||||
for _, tc := range validCases {
|
||||
t.Run(fmt.Sprintf("valid_%s", tc), func(t *testing.T) {
|
||||
result, err := validateIdentifier(tc, "test")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc, result)
|
||||
})
|
||||
}
|
||||
|
||||
// Test invalid identifiers
|
||||
invalidCases := []struct {
|
||||
name string
|
||||
value string
|
||||
}{
|
||||
{"empty", ""},
|
||||
{"space", "test table"},
|
||||
{"hyphen", "test-table"},
|
||||
{"dot", "test.table"},
|
||||
{"quote", `test"table`},
|
||||
{"semicolon", "test;table"},
|
||||
{"sql_injection", `test"; DROP KEYSPACE prod; --`},
|
||||
{"unicode", "test表"},
|
||||
}
|
||||
|
||||
for _, tc := range invalidCases {
|
||||
t.Run(fmt.Sprintf("invalid_%s", tc.name), func(t *testing.T) {
|
||||
_, err := validateIdentifier(tc.value, "test")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "invalid test name")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user