add context support to more storages

cassandra/cassandra_test.go

@@ -66,6 +66,22 @@ func Test_Set(t *testing.T) {
 	require.Equal(t, []byte("value"), val)
 }
 
+func Test_SetWithContext(t *testing.T) {
+	store := newTestStore(t)
+	defer store.Close()
+
+	// Test SetWithContext
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err := store.SetWithContext(ctx, "test", []byte("value"), 0)
+	require.ErrorIs(t, err, context.Canceled)
+
+	// Verify the value was not set
+	val, err := store.Get("test")
+	require.Error(t, err)
+	require.Empty(t, val)
+}
+
 // Test_Get tests the Get operation
 func Test_Get(t *testing.T) {
 	store := newTestStore(t)
@@ -86,6 +102,28 @@ func Test_Get(t *testing.T) {
 	require.Nil(t, val)
 }
 
+// Test_GetWithContext tests the Get operation with context
+func Test_GetWithContext(t *testing.T) {
+	store := newTestStore(t)
+	defer store.Close()
+
+	// Set a value first
+	err := store.Set("test", []byte("value"), 0)
+	require.NoError(t, err)
+
+	// Test GetWithContext
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	val, err := store.GetWithContext(ctx, "test")
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, val)
+
+	// Verify the value still exists
+	val, err = store.Get("test")
+	require.NoError(t, err)
+	require.Equal(t, []byte("value"), val)
+}
+
 // Test_Delete tests the Delete operation
 func Test_Delete(t *testing.T) {
 	store := newTestStore(t)
@@ -110,6 +148,32 @@ func Test_Delete(t *testing.T) {
 	require.Nil(t, val)
 }
 
+// Test_DeleteWithContext tests the Delete operation with context
+func Test_DeleteWithContext(t *testing.T) {
+	store := newTestStore(t)
+	defer store.Close()
+
+	// Set a value first
+	err := store.Set("test", []byte("value"), 0)
+	require.NoError(t, err)
+
+	// Verify the value exists
+	val, err := store.Get("test")
+	require.NoError(t, err)
+	require.Equal(t, []byte("value"), val)
+
+	// Test DeleteWithContext
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = store.DeleteWithContext(ctx, "test")
+	require.ErrorIs(t, err, context.Canceled)
+
+	// Verify the value still exists
+	val, err = store.Get("test")
+	require.NoError(t, err)
+	require.Equal(t, []byte("value"), val)
+}
+
 // Test_Expirable_Keys tests the expirable keys functionality
 func Test_Expirable_Keys(t *testing.T) {
 	store := newTestStore(t)
@@ -180,6 +244,33 @@ func Test_Reset(t *testing.T) {
 	require.Nil(t, val)
 }
 
+// Test_ResetWithContext tests the Reset method with context
+func Test_ResetWithContext(t *testing.T) {
+	store := newTestStore(t)
+	defer store.Close()
+
+	// Add some data
+	err := store.Set("test1", []byte("value1"), 0)
+	require.NoError(t, err)
+	err = store.Set("test2", []byte("value2"), 0)
+	require.NoError(t, err)
+
+	// Reset storage with context
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = store.ResetWithContext(ctx)
+	require.ErrorIs(t, err, context.Canceled)
+
+	// Verify data is still there
+	val, err := store.Get("test1")
+	require.NoError(t, err)
+	require.Equal(t, []byte("value1"), val)
+
+	val, err = store.Get("test2")
+	require.NoError(t, err)
+	require.Equal(t, []byte("value2"), val)
+}
+
 // Test_Valid_Identifiers tests valid identifier cases
 func Test_Valid_Identifiers(t *testing.T) {
 	store := newTestStore(t)