52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""测试视频生成次数重置功能"""
|
|||
|
|
from datetime import datetime, date
|
|||
|
|
from lover.db import SessionLocal
|
|||
|
|
from lover.models import User
|
|||
|
|
|
|||
|
|
def test_reset_logic():
|
|||
|
|
db = SessionLocal()
|
|||
|
|
try:
|
|||
|
|
user = db.query(User).filter(User.id == 70).with_for_update().first()
|
|||
|
|
if not user:
|
|||
|
|
print("用户不存在")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
print(f"重置前:")
|
|||
|
|
print(f" video_gen_remaining: {user.video_gen_remaining}")
|
|||
|
|
print(f" video_gen_reset_date: {user.video_gen_reset_date}")
|
|||
|
|
|
|||
|
|
# 模拟重置逻辑
|
|||
|
|
current_timestamp = int(datetime.utcnow().timestamp())
|
|||
|
|
is_vip = user.vip_endtime and user.vip_endtime > current_timestamp
|
|||
|
|
|
|||
|
|
if is_vip:
|
|||
|
|
last_reset = user.video_gen_reset_date
|
|||
|
|
today = datetime.utcnow().date()
|
|||
|
|
|
|||
|
|
print(f"\n检查:")
|
|||
|
|
print(f" 是否 VIP: {is_vip}")
|
|||
|
|
print(f" 上次重置日期: {last_reset}")
|
|||
|
|
print(f" 今天日期: {today}")
|
|||
|
|
print(f" 需要重置: {not last_reset or last_reset < today}")
|
|||
|
|
|
|||
|
|
if not last_reset or last_reset < today:
|
|||
|
|
user.video_gen_remaining = 2
|
|||
|
|
user.video_gen_reset_date = today
|
|||
|
|
db.add(user)
|
|||
|
|
db.commit()
|
|||
|
|
|
|||
|
|
print(f"\n重置后:")
|
|||
|
|
print(f" video_gen_remaining: {user.video_gen_remaining}")
|
|||
|
|
print(f" video_gen_reset_date: {user.video_gen_reset_date}")
|
|||
|
|
else:
|
|||
|
|
print("\n今天已经重置过了，不需要再次重置")
|
|||
|
|
else:
|
|||
|
|
print("用户不是 VIP，不重置")
|
|||
|
|
|
|||
|
|
finally:
|
|||
|
|
db.close()
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
test_reset_logic()
|