Skip to content

Commit 7b05314

Browse files
clnollclnoll
committed
fix: updates to Stigg calls (#17636)
1 parent 166464a commit 7b05314

File tree

4 files changed

+64
-66
lines changed

4 files changed

+64
-66
lines changed

‎airbyte-commons-entitlements/src/main/kotlin/io/airbyte/commons/entitlements/StiggCloudEntitlementClient.kt

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,8 @@ internal class StiggCloudEntitlementClient(
3636
organizationId: OrganizationId,
3737
plan: EntitlementPlan,
3838
) {
39-
val org =
40-
organizationService
41-
.getOrganization(organizationId.value)
42-
.orElseThrow {
43-
IllegalStateException("Organization $organizationId not found; could not add to plan $plan")
44-
} ?: throw IllegalStateException("getOrganization() returned null for $organizationId; could not add to plan $plan")
45-
4639
validatePlanChange(organizationId, plan)
47-
stigg.provisionCustomer(organizationId, org.name, plan)
40+
stigg.provisionCustomer(organizationId, plan)
4841

4942
logger.info { "Added organization to plan. organizationId=$organizationId plan=$plan" }
5043
}

‎airbyte-commons-entitlements/src/main/kotlin/io/airbyte/commons/entitlements/StiggEnterpriseEntitlementClient.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ class StiggEnterpriseEntitlementClient(
4040
.entitlements(
4141
OfflineEntitlements
4242
.builder()
43-
.customers(mapOf(smeOrgId.toString() to entitlements))
43+
.customers(mapOf(smeOrgId.value.toString() to entitlements))
4444
.build(),
4545
).build(),
4646
),

‎airbyte-commons-entitlements/src/main/kotlin/io/airbyte/commons/entitlements/StiggWrapper.kt

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
package io.airbyte.commons.entitlements
66

7+
import com.apollographql.apollo3.exception.ApolloException
78
import io.airbyte.commons.entitlements.models.Entitlement
89
import io.airbyte.commons.entitlements.models.EntitlementResult
910
import io.airbyte.domain.models.EntitlementPlan
@@ -17,6 +18,7 @@ import io.stigg.api.operations.type.ProvisionCustomerSubscriptionInput
1718
import io.stigg.sidecar.proto.v1.GetBooleanEntitlementRequest
1819
import io.stigg.sidecar.proto.v1.GetEntitlementsRequest
1920
import io.stigg.sidecar.sdk.Stigg
21+
import java.io.IOException
2022

2123
private val logger = KotlinLogging.logger {}
2224

@@ -35,32 +37,37 @@ internal class StiggWrapper(
3537
private val stigg: Stigg,
3638
) {
3739
fun getPlans(organizationId: OrganizationId): List<EntitlementPlan> {
38-
val resp =
39-
stigg.api().query(
40-
GetActiveSubscriptionsListQuery
41-
.builder()
42-
.input(
43-
GetActiveSubscriptionsInput
44-
.builder()
45-
.customerId(organizationId.toString())
46-
.build(),
47-
).build(),
48-
)
49-
50-
return resp.getActiveSubscriptions.map { EntitlementPlan.valueOf(it.slimSubscriptionFragmentV2.plan.planId) }.toList()
40+
try {
41+
val resp =
42+
stigg.api().query(
43+
GetActiveSubscriptionsListQuery
44+
.builder()
45+
.input(
46+
GetActiveSubscriptionsInput
47+
.builder()
48+
.customerId(organizationId.value.toString())
49+
.build(),
50+
).build(),
51+
)
52+
return resp.getActiveSubscriptions.map { EntitlementPlan.valueOf(it.slimSubscriptionFragmentV2.plan.planId) }.toList()
53+
} catch (e: ApolloException) {
54+
if (e.localizedMessage != null && e.localizedMessage!!.contains("Customer not found")) {
55+
logger.info { "No active subscriptions; organization not present in Stigg. organizationId=$organizationId" }
56+
return emptyList()
57+
}
58+
throw e
59+
}
5160
}
5261

5362
fun provisionCustomer(
5463
orgId: OrganizationId,
55-
orgName: String,
5664
plan: EntitlementPlan,
5765
) {
5866
stigg.api().mutation(
5967
ProvisionCustomerMutation(
6068
ProvisionCustomerInput
6169
.builder()
62-
.customerId(orgId.toString())
63-
.additionalMetaData(mapOf("name" to orgName))
70+
.customerId(orgId.value.toString())
6471
.subscriptionParams(
6572
ProvisionCustomerSubscriptionInput
6673
.builder()
@@ -81,7 +88,7 @@ internal class StiggWrapper(
8188
stigg.getBooleanEntitlement(
8289
GetBooleanEntitlementRequest
8390
.newBuilder()
84-
.setCustomerId(organizationId.toString())
91+
.setCustomerId(organizationId.value.toString())
8592
.setFeatureId(entitlement.featureId)
8693
.build(),
8794
)
@@ -105,7 +112,7 @@ internal class StiggWrapper(
105112
stigg.getEntitlements(
106113
GetEntitlementsRequest
107114
.newBuilder()
108-
.setCustomerId(organizationId.toString())
115+
.setCustomerId(organizationId.value.toString())
109116
.build(),
110117
)
111118

‎airbyte-commons-entitlements/src/test/kotlin/io/airbyte/commons/entitlements/StiggCloudEntitlementClientTest.kt

Lines changed: 37 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -80,77 +80,60 @@ internal class StiggCloudEntitlementClientTest {
8080
}
8181

8282
@Test
83-
fun `addOrganization with no org throws`() {
84-
val stigg = mockk<StiggWrapper>()
83+
fun `addOrganization calls validatePlanChange and provisionCustomer`() {
84+
val stigg = mockk<StiggWrapper>(relaxed = true)
8585
val client = StiggCloudEntitlementClient(stigg, orgService)
8686
val plan = EntitlementPlan.STANDARD
8787

88-
every { orgService.getOrganization(org1.value) } returns Optional.empty()
88+
every { stigg.getPlans(org1) } returns emptyList()
89+
90+
client.addOrganization(org1, plan)
8991

90-
assertThrows<IllegalStateException> { client.addOrganization(org1, plan) }
92+
verify { stigg.getPlans(org1) }
93+
verify { stigg.provisionCustomer(org1, plan) }
9194
}
9295

9396
@Test
9497
fun addOrganization() {
9598
val stigg = mockk<StiggWrapper>(relaxed = true)
9699
val client = StiggCloudEntitlementClient(stigg, orgService)
97-
val org =
98-
Organization()
99-
.withOrganizationId(org1.value)
100-
.withName("Org 1")
101100

102-
// mock organization
103-
every { orgService.getOrganization(org1.value) } returns Optional.of(org)
101+
every { stigg.getPlans(org1) } returns emptyList()
104102

105103
client.addOrganization(org1, EntitlementPlan.PRO)
106104

107-
verify { stigg.provisionCustomer(org1, org.name, EntitlementPlan.PRO) }
105+
verify { stigg.provisionCustomer(org1, EntitlementPlan.PRO) }
108106
}
109107

110108
@Test
111109
fun `upgrade from standard to pro`() {
112110
val stigg = mockk<StiggWrapper>(relaxed = true)
113111
val client = StiggCloudEntitlementClient(stigg, orgService)
114-
val org =
115-
Organization()
116-
.withOrganizationId(org1.value)
117-
.withName("Org 1")
118112

119-
every { orgService.getOrganization(org1.value) } returns Optional.of(org)
120113
every { stigg.getPlans(org1) } returns listOf(EntitlementPlan.STANDARD)
121114

122115
client.addOrganization(org1, EntitlementPlan.PRO)
123116

124-
verify { stigg.provisionCustomer(org1, org.name, EntitlementPlan.PRO) }
117+
verify { stigg.provisionCustomer(org1, EntitlementPlan.PRO) }
125118
}
126119

127120
@Test
128121
fun `adding back to same plan is a no-op`() {
129122
val stigg = mockk<StiggWrapper>(relaxed = true)
130123
val client = StiggCloudEntitlementClient(stigg, orgService)
131-
val org =
132-
Organization()
133-
.withOrganizationId(org1.value)
134-
.withName("Org 1")
135124

136-
every { orgService.getOrganization(org1.value) } returns Optional.of(org)
137125
every { stigg.getPlans(org1) } returns listOf(EntitlementPlan.STANDARD)
138126

139127
client.addOrganization(org1, EntitlementPlan.STANDARD)
140128

141-
verify { stigg.provisionCustomer(org1, org.name, EntitlementPlan.STANDARD) }
129+
verify { stigg.provisionCustomer(org1, EntitlementPlan.STANDARD) }
142130
}
143131

144132
@Test
145133
fun `cannot downgrade from pro to standard`() {
146134
val stigg = mockk<StiggWrapper>(relaxed = true)
147135
val client = StiggCloudEntitlementClient(stigg, orgService)
148-
val org =
149-
Organization()
150-
.withOrganizationId(org1.value)
151-
.withName("Org 1")
152136

153-
every { orgService.getOrganization(org1.value) } returns Optional.of(org)
154137
every { stigg.getPlans(org1) } returns listOf(EntitlementPlan.PRO)
155138

156139
val exception =
@@ -171,12 +154,7 @@ internal class StiggCloudEntitlementClientTest {
171154
fun `cannot downgrade from pro trial to standard`() {
172155
val stigg = mockk<StiggWrapper>(relaxed = true)
173156
val client = StiggCloudEntitlementClient(stigg, orgService)
174-
val org =
175-
Organization()
176-
.withOrganizationId(org1.value)
177-
.withName("Org 1")
178157

179-
every { orgService.getOrganization(org1.value) } returns Optional.of(org)
180158
every { stigg.getPlans(org1) } returns listOf(EntitlementPlan.PRO_TRIAL)
181159

182160
val exception =
@@ -197,12 +175,7 @@ internal class StiggCloudEntitlementClientTest {
197175
fun `plan validation handles multiple plans when checking for downgrades`() {
198176
val stigg = mockk<StiggWrapper>(relaxed = true)
199177
val client = StiggCloudEntitlementClient(stigg, orgService)
200-
val org =
201-
Organization()
202-
.withOrganizationId(org1.value)
203-
.withName("Org 1")
204178

205-
every { orgService.getOrganization(org1.value) } returns Optional.of(org)
206179
every { stigg.getPlans(org1) } returns listOf(EntitlementPlan.CORE, EntitlementPlan.PRO)
207180

208181
val exception =
@@ -219,6 +192,31 @@ internal class StiggCloudEntitlementClientTest {
219192
data.errorMessage,
220193
)
221194
}
195+
196+
@Test
197+
fun `addOrganization handles stigg wrapper exceptions gracefully`() {
198+
val stigg = mockk<StiggWrapper>(relaxed = true)
199+
val client = StiggCloudEntitlementClient(stigg, orgService)
200+
201+
every { stigg.getPlans(org1) } throws RuntimeException("Stigg service unavailable")
202+
203+
assertThrows<RuntimeException> {
204+
client.addOrganization(org1, EntitlementPlan.STANDARD)
205+
}
206+
}
207+
208+
@Test
209+
fun `addOrganization works with organization not found in stigg`() {
210+
val stigg = mockk<StiggWrapper>(relaxed = true)
211+
val client = StiggCloudEntitlementClient(stigg, orgService)
212+
213+
// Simulate organization not being found in Stigg (empty plans)
214+
every { stigg.getPlans(org1) } returns emptyList()
215+
216+
client.addOrganization(org1, EntitlementPlan.STANDARD)
217+
218+
verify { stigg.provisionCustomer(org1, EntitlementPlan.STANDARD) }
219+
}
222220
}
223221

224222
private fun EntitlementResult.assertEntitled() {
@@ -242,7 +240,7 @@ private fun buildOfflineClient(vararg entitlements: Pair<OrganizationId, String>
242240
.groupBy({ it.first }, { it.second })
243241
.entries
244242
.associate { entry ->
245-
entry.key.toString() to
243+
entry.key.value.toString() to
246244
CustomerEntitlements
247245
.builder()
248246
.entitlements(

0 commit comments

Comments
 (0)