generated from obsidianmd/obsidian-sample-plugin
-
-
Notifications
You must be signed in to change notification settings - Fork 61
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
c0e1986
commit 084479e
Showing
7 changed files
with
1,008 additions
and
705 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,119 @@ | ||
import { db, UserUsageTable } from "@/drizzle/schema"; | ||
import { PRODUCTS } from "@/srm.config"; | ||
import { eq } from "drizzle-orm"; | ||
import type { Server } from "http"; | ||
import { createServer } from "http"; | ||
import { NextApiHandler } from "next"; | ||
import { GET } from "./route"; | ||
import type { SuperTest, Test } from "supertest"; | ||
import supertest from "supertest"; | ||
/** | ||
* @jest-environment node | ||
*/ | ||
|
||
describe("Token Reset Cron Job", () => { | ||
const mockUserId = "test-user-123"; | ||
const monthlyTokenLimit = 5000 * 1000; // 5M tokens | ||
let server: Server; | ||
let request: SuperTest<Test>; | ||
|
||
beforeAll(() => { | ||
const handler: NextApiHandler = (req, res) => { | ||
if (req.method === "GET") { | ||
return GET(req as any); | ||
} | ||
}; | ||
server = createServer(handler as any); | ||
request = supertest(server); | ||
}); | ||
|
||
afterAll((done) => { | ||
server.close(done); | ||
}); | ||
|
||
beforeEach(async () => { | ||
// Setup test data | ||
await db.insert(UserUsageTable).values({ | ||
userId: mockUserId, | ||
subscriptionStatus: "active", | ||
paymentStatus: "paid", | ||
tokenUsage: 1000000, // 1M tokens used | ||
maxTokenUsage: monthlyTokenLimit, | ||
billingCycle: "subscription", | ||
currentPlan: PRODUCTS.SubscriptionMonthly.metadata.plan, | ||
}); | ||
}); | ||
|
||
afterEach(async () => { | ||
// Cleanup test data | ||
await db.delete(UserUsageTable).where(eq(UserUsageTable.userId, mockUserId)); | ||
}); | ||
|
||
it("should reset token usage for active subscribers", async () => { | ||
const response = await request | ||
.get("/api/cron/reset-tokens") | ||
.set("authorization", `Bearer ${process.env.CRON_SECRET}`); | ||
|
||
expect(response.status).toBe(200); | ||
expect(response.body).toEqual({ | ||
success: true, | ||
message: "Token usage reset successful", | ||
}); | ||
|
||
// Verify token usage was reset | ||
const userUsage = await db | ||
.select() | ||
.from(UserUsageTable) | ||
.where(eq(UserUsageTable.userId, mockUserId)); | ||
|
||
expect(userUsage[0].tokenUsage).toBe(0); | ||
expect(userUsage[0].maxTokenUsage).toBe(monthlyTokenLimit); | ||
}); | ||
|
||
it("should not reset tokens for inactive subscriptions", async () => { | ||
// Update user to inactive | ||
await db | ||
.update(UserUsageTable) | ||
.set({ subscriptionStatus: "inactive" }) | ||
.where(eq(UserUsageTable.userId, mockUserId)); | ||
|
||
const response = await request | ||
.get("/api/cron/reset-tokens") | ||
.set("authorization", `Bearer ${process.env.CRON_SECRET}`); | ||
|
||
expect(response.status).toBe(200); | ||
|
||
// Verify token usage was not reset | ||
const userUsage = await db | ||
.select() | ||
.from(UserUsageTable) | ||
.where(eq(UserUsageTable.userId, mockUserId)); | ||
|
||
expect(userUsage[0].tokenUsage).toBe(1000000); // Should remain unchanged | ||
}); | ||
|
||
it("should return 401 for unauthorized requests", async () => { | ||
const response = await request | ||
.get("/api/cron/reset-tokens") | ||
.set("authorization", "Bearer invalid-token"); | ||
|
||
expect(response.status).toBe(401); | ||
}); | ||
|
||
it("should handle database errors gracefully", async () => { | ||
// Mock a database error | ||
jest.spyOn(db, "update").mockRejectedValueOnce( | ||
new Error("Database error") as never | ||
); | ||
|
||
const response = await request | ||
.get("/api/cron/reset-tokens") | ||
.set("authorization", `Bearer ${process.env.CRON_SECRET}`); | ||
|
||
expect(response.status).toBe(500); | ||
expect(response.body).toEqual({ | ||
success: false, | ||
error: "Failed to reset token usage", | ||
}); | ||
}); | ||
}); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,45 @@ | ||
import { db, UserUsageTable } from "@/drizzle/schema"; | ||
import { PRODUCTS } from "@/srm.config"; | ||
import { eq, and } from "drizzle-orm"; | ||
import { NextResponse } from "next/server"; | ||
|
||
export const runtime = "edge"; | ||
|
||
async function resetTokenUsage() { | ||
const monthlyTokenLimit = 5000 * 1000; // 5M tokens | ||
|
||
// Reset tokens for active subscribers with valid plans | ||
await db | ||
.update(UserUsageTable) | ||
.set({ | ||
tokenUsage: 0, | ||
maxTokenUsage: monthlyTokenLimit, | ||
}) | ||
.where( | ||
and( | ||
eq(UserUsageTable.subscriptionStatus, "active"), | ||
eq(UserUsageTable.paymentStatus, "paid") | ||
) | ||
); | ||
|
||
return { success: true, message: "Token usage reset successful" }; | ||
} | ||
|
||
export async function GET(request: Request) { | ||
try { | ||
// Verify that the request is coming from Vercel Cron | ||
const authHeader = request.headers.get("authorization"); | ||
if (authHeader !== `Bearer ${process.env.CRON_SECRET}`) { | ||
return new NextResponse("Unauthorized", { status: 401 }); | ||
} | ||
|
||
const result = await resetTokenUsage(); | ||
return NextResponse.json(result); | ||
} catch (error) { | ||
console.error("Error resetting token usage:", error); | ||
return NextResponse.json( | ||
{ success: false, error: "Failed to reset token usage" }, | ||
{ status: 500 } | ||
); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.