158 lines
6 KiB
PL/PgSQL
158 lines
6 KiB
PL/PgSQL
\set ON_ERROR_STOP on
|
|
|
|
-- Regression: weekly quota consumption must spend overage credits once the
|
|
-- 7-day window total reaches the base limit, and usage from 7 days ago must
|
|
-- no longer count (window = CURRENT_DATE-6 .. CURRENT_DATE).
|
|
|
|
BEGIN;
|
|
|
|
CREATE OR REPLACE FUNCTION pg_temp.assert_true(condition boolean, message text)
|
|
RETURNS void
|
|
LANGUAGE plpgsql
|
|
AS $$
|
|
BEGIN
|
|
IF condition IS NOT TRUE THEN
|
|
RAISE EXCEPTION 'assertion_failed: %', message;
|
|
END IF;
|
|
END;
|
|
$$;
|
|
|
|
INSERT INTO auth.users (
|
|
id, aud, role, email, encrypted_password, email_confirmed_at,
|
|
raw_app_meta_data, raw_user_meta_data, created_at, updated_at
|
|
) VALUES
|
|
(
|
|
'31000000-0000-4000-8000-000000000001', 'authenticated', 'authenticated',
|
|
'consume-quota-weekly@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
|
|
'{"provider":"email","providers":["email"]}'::jsonb, '{}'::jsonb, now(), now()
|
|
),
|
|
(
|
|
'31000000-0000-4000-8000-000000000002', 'authenticated', 'authenticated',
|
|
'consume-quota-daily@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
|
|
'{"provider":"email","providers":["email"]}'::jsonb, '{}'::jsonb, now(), now()
|
|
);
|
|
|
|
UPDATE public.subscriptions SET tier = 'free', status = 'active', overage_credits = 1
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
|
UPDATE public.subscriptions SET tier = 'pro', status = 'active', overage_credits = 0
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000002';
|
|
|
|
SELECT pg_temp.assert_true(
|
|
NOT has_function_privilege('authenticated', 'public.consume_quota(uuid,text,integer,text)', 'EXECUTE'),
|
|
'authenticated users cannot consume quota directly'
|
|
);
|
|
SELECT pg_temp.assert_true(
|
|
has_function_privilege('service_role', 'public.consume_quota(uuid,text,integer,text)', 'EXECUTE'),
|
|
'service role can consume quota'
|
|
);
|
|
|
|
-- 250 Haiku calls spread over the six days before today (about 42 a day),
|
|
-- plus a large row 7 days ago that has already left the weekly window.
|
|
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
|
SELECT '31000000-0000-4000-8000-000000000001', CURRENT_DATE - d, 'llm_haiku',
|
|
CASE WHEN d = 1 THEN 250 - 41 * 5 ELSE 41 END
|
|
FROM generate_series(1, 6) AS d;
|
|
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
|
VALUES ('31000000-0000-4000-8000-000000000001', CURRENT_DATE - 7, 'llm_haiku', 200);
|
|
|
|
DO $$
|
|
DECLARE
|
|
first_call jsonb;
|
|
second_call jsonb;
|
|
credits integer;
|
|
today_count integer;
|
|
BEGIN
|
|
first_call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_haiku', 250, 'weekly');
|
|
SELECT overage_credits INTO credits FROM public.subscriptions
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
|
PERFORM pg_temp.assert_true(
|
|
(first_call->>'allowed')::boolean
|
|
AND first_call->>'consumed_from' = 'overage'
|
|
AND (first_call->>'overage_credits')::integer = 0
|
|
AND (first_call->>'current')::integer = 251
|
|
AND credits = 0,
|
|
format('weekly total at the limit spends the overage credit: %s credits=%s', first_call, credits)
|
|
);
|
|
|
|
second_call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_haiku', 250, 'weekly');
|
|
SELECT count INTO today_count FROM public.daily_usage
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000001' AND feature = 'llm_haiku' AND date = CURRENT_DATE;
|
|
PERFORM pg_temp.assert_true(
|
|
NOT (second_call->>'allowed')::boolean
|
|
AND second_call->>'consumed_from' = 'none'
|
|
AND today_count = 1,
|
|
format('no base and no credit left denies without counting: %s today=%s', second_call, today_count)
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
-- The row from 7 days ago rolls off: removing one day inside the window
|
|
-- frees base allowance even though CURRENT_DATE-7 still holds 200 uses.
|
|
DELETE FROM public.daily_usage
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000001' AND feature = 'llm_haiku' AND date = CURRENT_DATE - 2;
|
|
|
|
DO $$
|
|
DECLARE
|
|
call jsonb;
|
|
BEGIN
|
|
call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_haiku', 250, 'weekly');
|
|
PERFORM pg_temp.assert_true(
|
|
(call->>'allowed')::boolean AND call->>'consumed_from' = 'base',
|
|
format('usage older than CURRENT_DATE-6 does not count: %s', call)
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
-- Not-available features never spend credits.
|
|
UPDATE public.subscriptions SET overage_credits = 3
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
|
DO $$
|
|
DECLARE
|
|
call jsonb;
|
|
credits integer;
|
|
BEGIN
|
|
call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_sonnet', 0, 'daily');
|
|
SELECT overage_credits INTO credits FROM public.subscriptions
|
|
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
|
PERFORM pg_temp.assert_true(
|
|
NOT (call->>'allowed')::boolean AND call->>'consumed_from' = 'none' AND credits = 3,
|
|
format('limit 0 is denied without touching credits: %s credits=%s', call, credits)
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
-- Daily policies (and the legacy three-argument call) still count only today.
|
|
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
|
VALUES ('31000000-0000-4000-8000-000000000002', CURRENT_DATE - 1, 'llm_sonnet', 300);
|
|
|
|
DO $$
|
|
DECLARE
|
|
daily_call jsonb;
|
|
legacy_call jsonb;
|
|
unlimited_call jsonb;
|
|
BEGIN
|
|
daily_call := public.consume_quota('31000000-0000-4000-8000-000000000002', 'llm_sonnet', 300, 'daily');
|
|
legacy_call := public.consume_quota('31000000-0000-4000-8000-000000000002', 'llm_sonnet', 300);
|
|
unlimited_call := public.consume_quota('31000000-0000-4000-8000-000000000002', 'stt_transcribe', -1, 'daily');
|
|
PERFORM pg_temp.assert_true(
|
|
daily_call->>'consumed_from' = 'base'
|
|
AND (daily_call->>'current')::integer = 1
|
|
AND legacy_call->>'consumed_from' = 'base'
|
|
AND (legacy_call->>'current')::integer = 2
|
|
AND unlimited_call->>'consumed_from' = 'unlimited'
|
|
AND (unlimited_call->>'limit')::integer = -1,
|
|
format('daily window ignores yesterday: %s / %s / %s', daily_call, legacy_call, unlimited_call)
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
DO $$
|
|
BEGIN
|
|
PERFORM public.consume_quota('31000000-0000-4000-8000-000000000002', 'llm_sonnet', 300, 'monthly');
|
|
RAISE EXCEPTION 'assertion_failed: unknown period was accepted';
|
|
EXCEPTION
|
|
WHEN invalid_parameter_value THEN NULL;
|
|
END;
|
|
$$;
|
|
|
|
ROLLBACK;
|