prompt_payload_store_integration_test.go 3.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. package securityaudit
  2. import (
  3. "context"
  4. "os"
  5. "strings"
  6. "testing"
  7. "time"
  8. "github.com/redis/go-redis/v9"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestRedisPayloadStoreRoundTripTTLNamespaceAndDelete(t *testing.T) {
  12. address := strings.TrimSpace(os.Getenv(promptAuditRedisTestEnv))
  13. if address == "" {
  14. t.Skip(promptAuditRedisTestEnv + " is not set")
  15. }
  16. client := redis.NewClient(&redis.Options{Addr: address})
  17. t.Cleanup(func() { require.NoError(t, client.Close()) })
  18. store := NewRedisPayloadStore(client)
  19. ctx := context.Background()
  20. const jobID int64 = 987654321
  21. const canary = "PROMPT_CANARY_REDIS_ONLY_PAYLOAD"
  22. _ = store.Delete(ctx, jobID)
  23. require.NoError(t, store.Set(ctx, jobID, canary, 2*DefaultPayloadTTL))
  24. require.Equal(t, PayloadKeyPrefix+"987654321", payloadKey(jobID))
  25. value, err := store.Get(ctx, jobID)
  26. require.NoError(t, err)
  27. require.Equal(t, canary, value)
  28. ttl, err := client.TTL(ctx, payloadKey(jobID)).Result()
  29. require.NoError(t, err)
  30. require.Greater(t, ttl, time.Duration(0))
  31. require.LessOrEqual(t, ttl, DefaultPayloadTTL)
  32. require.NoError(t, store.Delete(ctx, jobID))
  33. _, err = store.Get(ctx, jobID)
  34. require.ErrorIs(t, err, redis.Nil)
  35. }
  36. func TestPromptRuntimeAggregatesConfigWorkersQueueRedisEndpointsAndGuardMetrics(t *testing.T) {
  37. address := strings.TrimSpace(os.Getenv(promptAuditRedisTestEnv))
  38. if address == "" {
  39. t.Skip(promptAuditRedisTestEnv + " is not set")
  40. }
  41. db := openPromptAuditIntegrationDB(t)
  42. client := redis.NewClient(&redis.Options{Addr: address})
  43. t.Cleanup(func() { require.NoError(t, client.Close()) })
  44. config := &fakeConfigStore{active: true, cfg: ActiveConfig{
  45. RiskControlEnabled: true, Enabled: true, WorkerCount: 3, QueueCapacity: 123,
  46. ConfigVersion: 9, AllGroups: true,
  47. }}
  48. metrics := NewAtomicMetrics()
  49. metrics.Observe(DecisionBlock, 25*time.Millisecond)
  50. metrics.IncFailover()
  51. metrics.IncEnqueued()
  52. metrics.IncDropped()
  53. service := NewPromptService(
  54. config,
  55. NewPostgreSQLRepository(db),
  56. NewRedisPayloadStore(client),
  57. NewOpenAICompatibleScanner(),
  58. metrics,
  59. )
  60. service.probes["guard-1"] = ProbeResult{OK: true, Status: "healthy", HTTPStatus: 200}
  61. runtime := service.Runtime(context.Background())
  62. require.Equal(t, ModeAsync, runtime.EffectiveMode)
  63. require.Equal(t, int64(9), runtime.ExpectedConfigVersion)
  64. require.Equal(t, int64(9), runtime.ActiveConfigVersion)
  65. require.Equal(t, 3, runtime.WorkerTotal)
  66. require.Equal(t, 123, runtime.QueueCapacity)
  67. require.Equal(t, "ok", runtime.DatabaseStatus)
  68. require.Equal(t, "ok", runtime.RedisStatus)
  69. require.Contains(t, runtime.Endpoints, "guard-1")
  70. require.Equal(t, int64(1), runtime.GuardMetrics.Total)
  71. require.Equal(t, int64(1), runtime.GuardMetrics.Blocked)
  72. require.Equal(t, int64(1), runtime.GuardMetrics.Failovers)
  73. require.Equal(t, int64(25), runtime.GuardMetrics.LatencyP95MS)
  74. require.Equal(t, int64(1), runtime.EnqueuedTotal)
  75. require.Equal(t, int64(1), runtime.DroppedTotal)
  76. // The runner has not been started in this integration test, so the honest
  77. // process status is degraded rather than a fabricated running heartbeat.
  78. require.Equal(t, "degraded", runtime.ProcessStatus)
  79. }