test: enhance plugin state test coverage and readability (#1349)

Signed-off-by: Kay Yan <[email protected]>

Author: yankay

pkg/epp/plugins/plugin_state.go

@@ -109,17 +109,23 @@ func (s *PluginState) cleanup(ctx context.Context) {
 			log.FromContext(ctx).V(logutil.DEFAULT).Info("Shutting down plugin state cleanup")
 			return
 		case <-ticker.C:
-			s.requestToLastAccessTime.Range(func(k, v any) bool {
-				requestID := k.(string)
-				lastAccessTime := v.(time.Time)
-				if time.Since(lastAccessTime) > stalenessThreshold {
-					s.Delete(requestID) // cleanup stale requests (this is safe in sync.Map)
-				}
-				return true
-			})
+			s.cleanStaleRequests()
 		}
 	}
+}
 
+// cleanStaleRequests iterates through all requests and removes those that haven't been
+// accessed for longer than stalenessThreshold. This operation is safe to run concurrently
+// with other operations on the PluginState.
+func (s *PluginState) cleanStaleRequests() {
+	s.requestToLastAccessTime.Range(func(k, v any) bool {
+		requestID := k.(string)
+		lastAccessTime := v.(time.Time)
+		if time.Since(lastAccessTime) > stalenessThreshold {
+			s.Delete(requestID) // cleanup stale requests (this is safe in sync.Map)
+		}
+		return true
+	})
 }
 
 // ReadPluginStateKey retrieves data with the given key from PluginState and asserts it to type T.

pkg/epp/plugins/plugin_state_test.go

@@ -0,0 +1,128 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package plugins
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+// pluginTestData implements the StateData interface for testing purposes.
+// It provides a simple string value that can be stored and retrieved.
+type pluginTestData struct {
+	value string
+}
+
+// Clone implements the StateData interface, creating a deep copy of the data.
+func (d *pluginTestData) Clone() StateData {
+	if d == nil {
+		return nil
+	}
+	return &pluginTestData{value: d.value}
+}
+
+// TestPluginState_ReadWrite verifies the basic operations of PluginState:
+// - Writing data for a request
+// - Reading the data back
+// - Deleting the data and confirming it's removed
+func TestPluginState_ReadWrite(t *testing.T) {
+	ctx := logutil.NewTestLoggerIntoContext(context.Background())
+
+	state := NewPluginState(ctx)
+
+	req1 := "req1"
+	key := StateKey("key")
+	data1 := "bar1"
+	req2 := "req2"
+	data2 := "bar2"
+
+	// Write data to the state storage
+	state.Write(req1, key, &pluginTestData{value: data1})
+	state.Write(req2, key, &pluginTestData{value: data2})
+
+	// Read back the req1 data and verify its content
+	readData, err := state.Read(req1, key)
+	assert.NoError(t, err)
+	td, ok := readData.(*pluginTestData)
+	assert.True(t, ok, "should be able to cast to pluginTestData")
+	assert.Equal(t, data1, td.value)
+
+	// Delete the req2 data and verify it's removed
+	state.Delete(req2)
+	readData, err = state.Read(req2, key)
+	assert.Equal(t, ErrNotFound, err)
+	assert.Nil(t, readData, "expected no data after delete")
+
+	// Read back the req1 data and verify its content after the req2 deleted
+	readData, err = state.Read(req1, key)
+	assert.NoError(t, err)
+	td, ok = readData.(*pluginTestData)
+	assert.True(t, ok, "should be able to cast to pluginTestData")
+	assert.Equal(t, data1, td.value)
+}
+
+// TestReadPluginStateKey tests the generic helper function ReadPluginStateKey which provides
+// type-safe access to stored data. It verifies:
+// - Successful type assertion and data retrieval
+// - Error handling for non-existent keys
+func TestReadPluginStateKey(t *testing.T) {
+	ctx := logutil.NewTestLoggerIntoContext(context.Background())
+	state := NewPluginState(ctx)
+
+	requestID := "req-1"
+	key := StateKey("foo")
+	data := &pluginTestData{value: "bar"}
+
+	state.Write(requestID, key, data)
+
+	// Read
+	val, err := ReadPluginStateKey[*pluginTestData](state, requestID, key)
+	assert.NoError(t, err)
+	assert.Equal(t, "bar", val.value)
+
+	// Not Found
+	_, err = ReadPluginStateKey[*pluginTestData](state, "not-exist", key)
+	assert.Equal(t, ErrNotFound, err)
+}
+
+// TestPluginState_Cleanup verifies the automatic cleanup of stale data.
+// It tests that data which hasn't been accessed for longer than stalenessThreshold
+// is properly removed from the storage.
+func TestPluginState_Cleanup(t *testing.T) {
+	ctx := logutil.NewTestLoggerIntoContext(context.Background())
+
+	state := NewPluginState(ctx)
+
+	requestID := "req-stale"
+	key := StateKey("foo")
+	data := &pluginTestData{value: "bar"}
+
+	state.Write(requestID, key, data)
+
+	// Manually set last access time to far in the past
+	state.requestToLastAccessTime.Store(requestID, time.Now().Add(-2*stalenessThreshold))
+	// Manually CleanUp
+	state.cleanStaleRequests()
+
+	_, err := state.Read(requestID, key)
+	assert.Equal(t, ErrNotFound, err)
+}