show expiration warning for trial accounts

rodrimaia · rodrimaia · commit 2c7c2ee15f9a · 2023-04-26T14:02:42.000Z
diff --git a/enterprise/coderd/license/license.go b/enterprise/coderd/license/license.go
@@ -74,7 +74,13 @@ func Entitlements(
 
 		// Add warning if license is expiring soon
 		daysToExpire := int(math.Ceil(claims.LicenseExpires.Sub(now).Hours() / 24))
-		if daysToExpire > 0 && daysToExpire < 30 {
+		isTrial := entitlements.Trial
+		showWarningDays := 30
+		if isTrial {
+			showWarningDays = 7
+		}
+		isExpiringSoon := daysToExpire > 0 && daysToExpire < showWarningDays
+		if isExpiringSoon {
 			day := "day"
 			if daysToExpire > 1 {
 				day = "days"
diff --git a/enterprise/coderd/license/license_test.go b/enterprise/coderd/license/license_test.go
@@ -130,6 +130,95 @@ func TestEntitlements(t *testing.T) {
 			"Your license expires in 2 days.",
 		)
 	})
+
+	t.Run("Expiration warning for license expiring in 1 day", func(t *testing.T) {
+		t.Parallel()
+		db := dbfake.New()
+		db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureUserLimit: 100,
+					codersdk.FeatureAuditLog:  1,
+				},
+
+				GraceAt:   time.Now().AddDate(0, 0, 1),
+				ExpiresAt: time.Now().AddDate(0, 0, 5),
+			}),
+			Exp: time.Now().AddDate(0, 0, 5),
+		})
+
+		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+
+		require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[codersdk.FeatureAuditLog].Entitlement)
+		require.Contains(
+			t, entitlements.Warnings,
+			"Your license expires in 1 day.",
+		)
+	})
+
+	t.Run("Expiration warning for trials", func(t *testing.T) {
+		t.Parallel()
+		db := dbfake.New()
+		db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureUserLimit: 100,
+					codersdk.FeatureAuditLog:  1,
+				},
+
+				Trial:     true,
+				GraceAt:   time.Now().AddDate(0, 0, 8),
+				ExpiresAt: time.Now().AddDate(0, 0, 5),
+			}),
+			Exp: time.Now().AddDate(0, 0, 5),
+		})
+
+		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.True(t, entitlements.Trial)
+
+		require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[codersdk.FeatureAuditLog].Entitlement)
+		require.NotContains( // it should not contain a warning since it is a trial license
+			t, entitlements.Warnings,
+			"Your license expires in 8 days.",
+		)
+	})
+
+	t.Run("Expiration warning for non trials", func(t *testing.T) {
+		t.Parallel()
+		db := dbfake.New()
+		db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureUserLimit: 100,
+					codersdk.FeatureAuditLog:  1,
+				},
+
+				GraceAt:   time.Now().AddDate(0, 0, 30),
+				ExpiresAt: time.Now().AddDate(0, 0, 5),
+			}),
+			Exp: time.Now().AddDate(0, 0, 5),
+		})
+
+		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+
+		require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[codersdk.FeatureAuditLog].Entitlement)
+		require.NotContains( // it should not contain a warning since it is a trial license
+			t, entitlements.Warnings,
+			"Your license expires in 30 days.",
+		)
+	})
+
 	t.Run("SingleLicenseNotEntitled", func(t *testing.T) {
 		t.Parallel()
 		db := dbfake.New()
