fix: Skip trial button visibility (#23034)

* fix: flaky e2e

* fix: skip trial button visibility

* fix type error

* fix test

* update

* Update skipTeamTrials.test.ts

* fix test

packages/features/shell/useBottomNavItems.ts

@@ -1,7 +1,7 @@
 import type { User as UserAuth } from "next-auth";
 
 import { IS_DUB_REFERRALS_ENABLED } from "@calcom/lib/constants";
-import { useHasActiveTeamPlan } from "@calcom/lib/hooks/useHasPaidPlan";
+import { useHasActiveTeamPlanAsOwner } from "@calcom/lib/hooks/useHasPaidPlan";
 import { useLocale } from "@calcom/lib/hooks/useLocale";
 import { trpc } from "@calcom/trpc/react";
 import { showToast } from "@calcom/ui/components/toast";
@@ -20,7 +20,7 @@ export function useBottomNavItems({
   user,
 }: BottomNavItemsProps): NavigationItemType[] {
   const { t } = useLocale();
-  const { isTrial } = useHasActiveTeamPlan();
+  const { isTrial } = useHasActiveTeamPlanAsOwner();
   const utils = trpc.useUtils();
 
   const skipTeamTrialsMutation = trpc.viewer.teams.skipTeamTrials.useMutation({

packages/lib/hooks/useHasPaidPlan.ts

@@ -46,5 +46,12 @@ export function useHasActiveTeamPlan() {
 
   return { isPending, hasActiveTeamPlan: !!data?.isActive, isTrial: !!data?.isTrial };
 }
+export function useHasActiveTeamPlanAsOwner() {
+  if (IS_SELF_HOSTED) return { isPending: false, hasActiveTeamPlan: true, isTrial: false };
+
+  const { data, isPending } = trpc.viewer.teams.hasActiveTeamPlan.useQuery({ ownerOnly: true });
+
+  return { isPending, hasActiveTeamPlan: !!data?.isActive, isTrial: !!data?.isTrial };
+}
 
 export default useHasPaidPlan;

packages/lib/server/repository/membership.ts

@@ -461,6 +461,21 @@ export class MembershipRepository {
     })) as unknown as Promise<MembershipDTOFromSelect<TSelect>[]>;
   }
 
+  static async findAllAcceptedTeamMemberships(userId: number, where?: Prisma.MembershipWhereInput) {
+    const teams = await prisma.team.findMany({
+      where: {
+        members: {
+          some: {
+            userId,
+            accepted: true,
+            ...(where ?? {}),
+          },
+        },
+      },
+    });
+    return teams;
+  }
+
   async findTeamAdminsByTeamId({ teamId }: { teamId: number }) {
     return await this.prismaClient.membership.findMany({
       where: {

packages/trpc/server/routers/viewer/teams/_router.tsx

@@ -13,6 +13,7 @@ import { ZGetInternalNotesPresetsInputSchema } from "./getInternalNotesPresets.s
 import { ZGetMemberAvailabilityInputSchema } from "./getMemberAvailability.schema";
 import { ZGetMembershipbyUserInputSchema } from "./getMembershipbyUser.schema";
 import { ZGetUserConnectedAppsInputSchema } from "./getUserConnectedApps.schema";
+import { ZHasActiveTeamPlanInputSchema } from "./hasActiveTeamPlan.schema";
 import { ZHasEditPermissionForUserSchema } from "./hasEditPermissionForUser.schema";
 import { ZInviteMemberInputSchema } from "./inviteMember/inviteMember.schema";
 import { ZInviteMemberByTokenSchemaInputSchema } from "./inviteMemberByToken.schema";
@@ -187,7 +188,7 @@ export const viewerTeamsRouter = router({
       const { default: handler } = await import("./updateInternalNotesPresets.handler");
       return handler({ ctx, input });
     }),
-  hasActiveTeamPlan: authedProcedure.query(async (opts) => {
+  hasActiveTeamPlan: authedProcedure.input(ZHasActiveTeamPlanInputSchema).query(async (opts) => {
     const { default: handler } = await import("./hasActiveTeamPlan.handler");
     return handler(opts);
   }),

packages/trpc/server/routers/viewer/teams/hasActiveTeamPlan.handler.ts

@@ -1,27 +1,30 @@
+import type { Prisma } from "@prisma/client";
+
 import { InternalTeamBilling } from "@calcom/ee/billing/teams/internal-team-billing";
 import { IS_SELF_HOSTED } from "@calcom/lib/constants";
+import { MembershipRepository } from "@calcom/lib/server/repository/membership";
 import { prisma } from "@calcom/prisma";
 import type { TrpcSessionUser } from "@calcom/trpc/server/types";
 
+import type { THasActiveTeamPlanInputSchema } from "./hasActiveTeamPlan.schema";
+
 type HasActiveTeamPlanOptions = {
   ctx: {
     user: Pick<NonNullable<TrpcSessionUser>, "id">;
   };
+  input: THasActiveTeamPlanInputSchema;
 };
 
-export const hasActiveTeamPlanHandler = async ({ ctx }: HasActiveTeamPlanOptions) => {
+export const hasActiveTeamPlanHandler = async ({ ctx, input }: HasActiveTeamPlanOptions) => {
   if (IS_SELF_HOSTED) return { isActive: true, isTrial: false };
 
-  const teams = await prisma.team.findMany({
-    where: {
-      members: {
-        some: {
-          userId: ctx.user.id,
-          accepted: true,
-        },
-      },
-    },
-  });
+  const whereClause: Prisma.MembershipWhereInput = { userId: ctx.user.id, accepted: true };
+
+  if (input?.ownerOnly) {
+    whereClause.role = "OWNER";
+  }
+
+  const teams = await MembershipRepository.findAllAcceptedTeamMemberships(ctx.user.id, whereClause);
 
   if (!teams.length) return { isActive: false, isTrial: false };
 

packages/trpc/server/routers/viewer/teams/hasActiveTeamPlan.schema.ts

@@ -1 +1,9 @@
-export {};
+import { z } from "zod";
+
+export const ZHasActiveTeamPlanInputSchema = z
+  .object({
+    ownerOnly: z.boolean().optional(),
+  })
+  .optional();
+
+export type THasActiveTeamPlanInputSchema = z.infer<typeof ZHasActiveTeamPlanInputSchema>;

packages/trpc/server/routers/viewer/teams/skipTeamTrials.handler.ts

@@ -1,6 +1,7 @@
 import { InternalTeamBilling } from "@calcom/ee/billing/teams/internal-team-billing";
 import { IS_SELF_HOSTED } from "@calcom/lib/constants";
 import logger from "@calcom/lib/logger";
+import { MembershipRepository } from "@calcom/lib/server/repository/membership";
 import { prisma } from "@calcom/prisma";
 import type { TrpcSessionUser } from "@calcom/trpc/server/types";
 
@@ -29,16 +30,8 @@ export const skipTeamTrialsHandler = async ({ ctx }: SkipTeamTrialsOptions) => {
       },
     });
 
-    const ownedTeams = await prisma.team.findMany({
-      where: {
-        members: {
-          some: {
-            userId: ctx.user.id,
-            accepted: true,
-            role: "OWNER",
-          },
-        },
-      },
+    const ownedTeams = await MembershipRepository.findAllAcceptedTeamMemberships(ctx.user.id, {
+      role: "OWNER",
     });
 
     for (const team of ownedTeams) {

packages/trpc/server/routers/viewer/teams/skipTeamTrials.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it, vi, beforeEach } from "vitest";
 
 import { InternalTeamBilling } from "@calcom/ee/billing/teams/internal-team-billing";
+import { MembershipRepository } from "@calcom/lib/server/repository/membership";
 import { prisma } from "@calcom/prisma";
 
 import { skipTeamTrialsHandler } from "./skipTeamTrials.handler";
@@ -31,9 +32,6 @@ vi.mock("@calcom/prisma", () => ({
     user: {
       update: vi.fn().mockResolvedValue({}),
     },
-    team: {
-      findMany: vi.fn(),
-    },
   },
 }));
 
@@ -46,6 +44,12 @@ vi.mock("@calcom/lib/logger", () => ({
   },
 }));
 
+vi.mock("@calcom/lib/server/repository/membership", () => ({
+  MembershipRepository: {
+    findAllAcceptedTeamMemberships: vi.fn(),
+  },
+}));
+
 const mockGetSubscriptionStatus = vi.fn();
 const mockEndTrial = vi.fn().mockResolvedValue(true);
 
@@ -69,7 +73,7 @@ describe("skipTeamTrialsHandler", () => {
   });
 
   it("should set user's trialEndsAt to null", async () => {
-    vi.mocked(prisma.team.findMany).mockResolvedValueOnce([]);
+    vi.mocked(MembershipRepository.findAllAcceptedTeamMemberships).mockResolvedValueOnce([]);
 
     // @ts-expect-error - simplified context for testing
     await skipTeamTrialsHandler({ ctx: mockCtx, input: {} });
@@ -89,9 +93,9 @@ describe("skipTeamTrialsHandler", () => {
     const mockTeams = [
       { id: 101, name: "Team 1" },
       { id: 102, name: "Team 2" },
-    ];
+    ] as any;
 
-    vi.mocked(prisma.team.findMany).mockResolvedValueOnce(mockTeams);
+    vi.mocked(MembershipRepository.findAllAcceptedTeamMemberships).mockResolvedValueOnce(mockTeams);
 
     mockGetSubscriptionStatus
       .mockResolvedValueOnce("trialing") // First team is in trial
@@ -102,16 +106,8 @@ describe("skipTeamTrialsHandler", () => {
 
     expect(prisma.user.update).toHaveBeenCalled();
 
-    expect(prisma.team.findMany).toHaveBeenCalledWith({
-      where: {
-        members: {
-          some: {
-            userId: mockCtx.user.id,
-            accepted: true,
-            role: "OWNER",
-          },
-        },
-      },
+    expect(MembershipRepository.findAllAcceptedTeamMemberships).toHaveBeenCalledWith(mockCtx.user.id, {
+      role: "OWNER",
     });
 
     expect(InternalTeamBilling).toHaveBeenCalledTimes(2);
@@ -119,7 +115,6 @@ describe("skipTeamTrialsHandler", () => {
     expect(InternalTeamBilling).toHaveBeenNthCalledWith(2, mockTeams[1]);
 
     expect(mockGetSubscriptionStatus).toHaveBeenCalledTimes(2);
-
     expect(mockEndTrial).toHaveBeenCalledTimes(1);
   });
 

packages/trpc/server/routers/viewer/workflows/update.handler.ts

@@ -92,7 +92,7 @@ export const updateHandler = async ({ ctx, input }: UpdateOptions) => {
 
   let teamsPlan = { isActive: false, isTrial: false };
   if (!isCurrentUsernamePremium) {
-    teamsPlan = await hasActiveTeamPlanHandler({ ctx });
+    teamsPlan = await hasActiveTeamPlanHandler({ ctx, input: { ownerOnly: false } });
   }
   const hasPaidPlan = IS_SELF_HOSTED || isCurrentUsernamePremium || teamsPlan.isActive;
   let newActiveOn: number[] = [];