easyai-ai-gateway/apps/api/internal/securityevents/metrics_test.go

package securityevents

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type metricsSnapshot struct {
	mode string
	last time.Time
	pool store.PostgresPoolMetricsSnapshot
}

func (m metricsSnapshot) SecurityEventMetrics(context.Context, string, string) (string, *time.Time, error) {
	return m.mode, &m.last, nil
}

func (m metricsSnapshot) PostgresPoolMetrics() store.PostgresPoolMetricsSnapshot {
	return m.pool
}

func (m metricsSnapshot) ListModelRateLimitStatuses(context.Context) ([]store.ModelRateLimitStatus, error) {
	return []store.ModelRateLimitStatus{{
		PlatformModelID: "platform-model-1",
		RPM:             store.RateLimitMetricStatus{Ratio: .25},
		TPM:             store.RateLimitMetricStatus{Ratio: .5},
		Concurrent:      store.RateLimitMetricStatus{Ratio: .75},
	}}, nil
}

func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
	metrics := &Metrics{}
	metrics.ObserveReceipt("accepted", 20*time.Millisecond, 2)
	metrics.ObserveReceipt("duplicate", 10*time.Millisecond, 0)
	metrics.ObserveWatermarkRejection()
	metrics.ObserveHeartbeat("accepted")
	metrics.ObserveIntrospection("failed")
	metrics.ObserveJWKSRefreshFailure("ssf")
	metrics.SetAsyncWorkerCapacity(96, 128, 96, true)
	metrics.SetWorkerLoad(12, 3, 8, 2, 5, 1, "busy")
	metrics.ObserveProviderQuotaWait("rpm")
	metrics.ObserveProviderQuotaWait("tpm_total")
	metrics.ObserveCandidateRouting("normal_rotation")
	metrics.ObserveCandidateRouting("full_avoided")
	metrics.ObserveCandidateRouting("quota_race_rotated")
	metrics.ObserveCandidateRouting("cooldown_skipped")
	metrics.ObserveCandidateRouting("disabled_skipped")
	metrics.ObserveCandidateRouting("all_full_queued")
	metrics.ObserveObjectStorage("write_success", "aliyun_oss", 128, 25*time.Millisecond)
	metrics.ObserveObjectStorage("write_failure", "s3", 0, 10*time.Millisecond)
	metrics.ObserveObjectStorage("retry", "s3", 0, 0)
	metrics.ObserveObjectStorage("failover", "s3", 0, 0)
	metrics.ObserveObjectStorage("all_failed", "", 0, 0)
	metrics.ObserveResultStorage("upstream_url")
	metrics.ObserveResultStorage("uploaded")
	metrics.ObserveResultDelivery("poll_success", 512)
	metrics.ObserveResultDelivery("poll_oversized", 70<<10)
	metrics.ObserveAsyncWorkerResize("success")
	metrics.ObserveConcurrencyLeaseRenewal("success")
	metrics.ObserveConcurrencyLeaseRenewal("lost")
	metrics.ObserveTaskEventSkip("duplicate")
	metrics.ObserveTaskEventSkip("budget_exceeded")
	metrics.ObserveAsyncAdmissionDispatch("admitted")
	metrics.ObserveAsyncAdmissionDispatch("global_wait")
	metrics.ObserveAsyncAdmissionDispatch("reselect")
	metrics.ObserveAsyncAdmissionDispatch("deferred")
	metrics.ObserveAsyncAdmissionDispatch("terminal")
	metrics.ObserveAsyncAdmissionDispatch("task_error")
	metrics.ObserveAsyncAdmissionDispatch("system_error")

	recorder := httptest.NewRecorder()
	metrics.Handler(metricsSnapshot{
		mode: "push_healthy",
		last: time.Now().Add(-time.Minute),
		pool: store.PostgresPoolMetricsSnapshot{
			MaxConnections: 32, TotalConnections: 18, AcquiredConnections: 12,
			IdleConnections: 6, EmptyAcquireCount: 3, CanceledAcquireCount: 1,
		},
	}, "issuer", "audience", true).
		ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/metrics", nil))
	if recorder.Code != http.StatusOK {
		t.Fatalf("metrics status=%d", recorder.Code)
	}
	for _, expected := range []string{
		`easyai_gateway_ssf_receipts_total{outcome="accepted"} 1`,
		`easyai_gateway_ssf_receipts_total{outcome="duplicate"} 1`,
		`easyai_gateway_ssf_sessions_deleted_total 2`,
		`easyai_gateway_ssf_watermark_rejections_total 1`,
		`easyai_gateway_ssf_processing_duration_seconds_bucket{le="0.05"} 2`,
		`easyai_gateway_ssf_mode{mode="push_healthy"} 1`,
		`easyai_gateway_oidc_introspection_total{outcome="failed"} 1`,
		`easyai_gateway_jwks_refresh_failures_total{outcome="ssf"} 1`,
		`easyai_gateway_async_worker_capacity 96`,
		`easyai_gateway_async_worker_desired_capacity 128`,
		`easyai_gateway_async_worker_capacity_capped 1`,
		`easyai_gateway_worker_safe_capacity 12`,
		`easyai_gateway_worker_heavy_capacity 3`,
		`easyai_gateway_worker_active_tasks 8`,
		`easyai_gateway_worker_preparing_tasks 2`,
		`easyai_gateway_worker_waiting_upstream_tasks 5`,
		`easyai_gateway_worker_finalizing_tasks 1`,
		`easyai_gateway_worker_pressure_state 1`,
		`easyai_gateway_provider_quota_waits_total{outcome="rpm"} 1`,
		`easyai_gateway_provider_quota_waits_total{outcome="tpm"} 1`,
		`easyai_gateway_candidate_routing_total{outcome="normal_rotation"} 1`,
		`easyai_gateway_candidate_routing_total{outcome="full_avoided"} 1`,
		`easyai_gateway_candidate_routing_total{outcome="quota_race_rotated"} 1`,
		`easyai_gateway_candidate_routing_total{outcome="cooldown_skipped"} 1`,
		`easyai_gateway_candidate_routing_total{outcome="disabled_skipped"} 1`,
		`easyai_gateway_candidate_routing_total{outcome="all_full_queued"} 1`,
		`easyai_gateway_storage_write_attempts_total{protocol="aliyun_oss"} 1`,
		`easyai_gateway_storage_write_failures_total{protocol="s3"} 1`,
		`easyai_gateway_storage_write_duration_seconds_total{protocol="aliyun_oss"} 0.025000`,
		`easyai_gateway_storage_channel_retries_total 1`,
		`easyai_gateway_storage_channel_failovers_total 1`,
		`easyai_gateway_storage_all_channels_failed_total 1`,
		`easyai_gateway_storage_objectified_bytes_total 128`,
		`easyai_gateway_result_storage_total{outcome="upstream_url"} 1`,
		`easyai_gateway_result_storage_total{outcome="uploaded"} 1`,
		`easyai_gateway_result_poll_responses_total 1`,
		`easyai_gateway_result_poll_response_bytes_total 512`,
		`easyai_gateway_result_poll_oversized_total 1`,
		`easyai_gateway_platform_model_rate_limit_utilization{platform_model_id="platform-model-1",metric="concurrent"} 0.750000`,
		`easyai_gateway_async_worker_resizes_total{outcome="success"} 1`,
		`easyai_gateway_concurrency_lease_renewals_total{outcome="success"} 1`,
		`easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`,
		`easyai_gateway_task_events_skipped_total{outcome="duplicate"} 1`,
		`easyai_gateway_task_events_skipped_total{outcome="budget_exceeded"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="admitted"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="global_wait"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="reselect"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="deferred"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="terminal"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="task_error"} 1`,
		`easyai_gateway_async_admission_dispatch_total{outcome="system_error"} 1`,
		`easyai_gateway_postgres_pool_max_connections 32`,
		`easyai_gateway_postgres_pool_total_connections 18`,
		`easyai_gateway_postgres_pool_acquired_connections 12`,
		`easyai_gateway_postgres_pool_idle_connections 6`,
		`easyai_gateway_postgres_pool_empty_acquire_total 3`,
		`easyai_gateway_postgres_pool_canceled_acquire_total 1`,
	} {
		if !strings.Contains(recorder.Body.String(), expected) {
			t.Fatalf("missing metric %q in:\n%s", expected, recorder.Body.String())
		}
	}
}