Skip to content

Commit f5b500b

Browse files
gosusnpgosusnp
committed
chore: refactor workload acceptance tests to use dynamic feature flags (#13492)
1 parent 0993ad0 commit f5b500b

File tree

9 files changed

+297
-28
lines changed

9 files changed

+297
-28
lines changed

‎airbyte-featureflag-server/build.gradle.kts

+1
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ dependencies {
2626
implementation(libs.jackson.databind)
2727
implementation(libs.jackson.dataformat)
2828
implementation(libs.jackson.kotlin)
29+
implementation(libs.kotlin.logging)
2930

3031
implementation(project(":oss:airbyte-commons"))
3132

‎airbyte-featureflag-server/src/main/kotlin/io/airbyte/featureflag/server/FeatureFlagService.kt

+20
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,19 @@ package io.airbyte.featureflag.server
33
import com.fasterxml.jackson.databind.ObjectMapper
44
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory
55
import com.fasterxml.jackson.module.kotlin.registerKotlinModule
6+
import io.airbyte.commons.json.Jsons
67
import io.airbyte.featureflag.server.model.Context
78
import io.airbyte.featureflag.server.model.FeatureFlag
89
import io.airbyte.featureflag.server.model.Rule
10+
import io.github.oshai.kotlinlogging.KotlinLogging
911
import io.micronaut.context.annotation.Property
1012
import jakarta.inject.Singleton
1113
import java.nio.file.Path
1214
import kotlin.io.path.exists
1315
import kotlin.io.path.isRegularFile
1416

17+
private val logger = KotlinLogging.logger {}
18+
1519
// This is open for testing, creating an interface might be the way to go
1620
@Singleton
1721
open class FeatureFlagService(
@@ -29,6 +33,7 @@ open class FeatureFlagService(
2933
}
3034
}
3135
}
36+
logger.info { "FeatureFlagService loaded with ${flags.toPrettyJson()}" }
3237
}
3338

3439
open fun delete(key: String) {
@@ -38,6 +43,15 @@ open class FeatureFlagService(
3843
open fun eval(
3944
key: String,
4045
context: Map<String, String>,
46+
): String? {
47+
val result = doEval(key, context)
48+
logger.debug { "Evaluating $key with $context to $result" }
49+
return result
50+
}
51+
52+
private fun doEval(
53+
key: String,
54+
context: Map<String, String>,
4155
): String? {
4256
val flag = flags[key] ?: return null
4357
for (rule in flag.rules) {
@@ -62,6 +76,7 @@ open class FeatureFlagService(
6276
throw Exception("$key already has a rule for context ${rule.context}")
6377
}
6478
flag.rules.add(rule.toMutableRule())
79+
logger.debug { "Updated $key to $flag" }
6580
return flag.toFeatureFlag()
6681
}
6782

@@ -74,6 +89,7 @@ open class FeatureFlagService(
7489
.find { it.context == rule.context }
7590
?.apply { value = rule.value }
7691
?: throw Exception("$key does not have a rule for context ${rule.context}")
92+
logger.debug { "Updated $key to $flag" }
7793
return flag.toFeatureFlag()
7894
}
7995

@@ -83,11 +99,13 @@ open class FeatureFlagService(
8399
): FeatureFlag {
84100
val flag = flags[key] ?: throw Exception("$key not found")
85101
flag.rules.removeIf { it.context == context }
102+
logger.debug { "Updated $key to $flag" }
86103
return flag.toFeatureFlag()
87104
}
88105

89106
open fun put(flag: FeatureFlag): FeatureFlag {
90107
flags[flag.key] = flag.toMutableFeatureFlag()
108+
logger.debug { "Updated ${flag.key} to $flag" }
91109
return get(flag.key) ?: throw Exception("Failed to put flag $flag")
92110
}
93111

@@ -100,6 +118,8 @@ open class FeatureFlagService(
100118
return put(flag)
101119
}
102120

121+
private fun <T : Any> T.toPrettyJson(): String = Jsons.toPrettyString(Jsons.jsonNode(this))
122+
103123
private fun Context.matches(env: Map<String, String>): Boolean = env[kind] == value
104124

105125
private fun MutableFeatureFlag.toFeatureFlag(): FeatureFlag = FeatureFlag(key = key, default = default, rules = rules.map { it.toRule() }.toList())

‎airbyte-featureflag/src/main/kotlin/tests/TestFlagsSetter.kt

+83-21
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,14 @@ import okhttp3.OkHttpClient
88
import okhttp3.Request
99
import okhttp3.RequestBody.Companion.toRequestBody
1010

11-
class TestFlagsSetter {
12-
private val baseurl = "http://local.airbyte.dev/api/v1/feature-flags"
11+
class TestFlagsSetter(baseUrl: String) {
12+
private val basePath = "/api/v1/feature-flags"
1313
private val httpClient = OkHttpClient().newBuilder().build()
14+
private val urlPrefix = if (baseUrl.endsWith("/")) "${baseUrl.trimEnd('/')}$basePath" else "$baseUrl$basePath"
1415

1516
class FlagOverride<T>(
1617
private val flag: Flag<T>,
17-
context: Context,
18+
context: Context? = null,
1819
value: T,
1920
private val testFlags: TestFlagsSetter,
2021
) : AutoCloseable {
@@ -27,61 +28,122 @@ class TestFlagsSetter {
2728
}
2829
}
2930

31+
class FlagRuleOverride<T>(
32+
private val flag: Flag<T>,
33+
private val context: Context,
34+
private val value: T,
35+
private val testFlags: TestFlagsSetter,
36+
) : AutoCloseable {
37+
init {
38+
testFlags.setRule(flag, context, value)
39+
}
40+
41+
override fun close() {
42+
testFlags.deleteRule(flag, context)
43+
}
44+
}
45+
3046
fun <T> withFlag(
3147
flag: Flag<T>,
32-
context: Context,
3348
value: T,
49+
context: Context? = null,
3450
) = FlagOverride(flag, context, value, this)
3551

3652
fun <T> deleteFlag(flag: Flag<T>) {
3753
httpClient.newCall(
3854
Request.Builder()
39-
.url("$baseurl/${flag.key}")
55+
.url("$urlPrefix/${flag.key}")
4056
.delete()
4157
.build(),
4258
).execute()
4359
}
4460

45-
fun <T> setFlag(
61+
fun <T> withRule(
4662
flag: Flag<T>,
4763
context: Context,
4864
value: T,
65+
) = FlagRuleOverride(flag, context, value, this)
66+
67+
fun <T> setFlag(
68+
flag: Flag<T>,
69+
context: Context? = null,
70+
value: T,
4971
) {
5072
val requestFlag =
5173
ApiFeatureFlag(
5274
key = flag.key,
5375
default = flag.default.toString(),
5476
rules =
55-
listOf(
56-
ApiRule(
57-
context = ApiContext(kind = context.kind, value = context.key),
58-
value = value.toString(),
59-
),
60-
),
77+
if (context != null) {
78+
listOf(
79+
ApiRule(
80+
context = ApiContext(kind = context.kind, value = context.key),
81+
value = value.toString(),
82+
),
83+
)
84+
} else {
85+
emptyList()
86+
},
6187
)
62-
httpClient.newCall(
88+
val response =
89+
httpClient.newCall(
90+
Request.Builder()
91+
.url(urlPrefix)
92+
.put(Jsons.serialize(requestFlag).toRequestBody("application/json".toMediaType()))
93+
.build(),
94+
).execute()
95+
assert(response.code == 200, { "Failed to update the feature flag ${requestFlag.key}, error: ${response.code}: ${response.body?.string()}" })
96+
}
97+
98+
fun <T> getFlag(flag: Flag<T>): String? {
99+
return httpClient.newCall(
63100
Request.Builder()
64-
.url(baseurl)
65-
.put(Jsons.serialize(requestFlag).toRequestBody("application/json".toMediaType()))
101+
.url("$urlPrefix/${flag.key}")
66102
.build(),
67103
).execute()
104+
.body?.string()
68105
}
69106

70-
fun <T> getFlag(flag: Flag<T>) {
71-
httpClient.newCall(
107+
fun <T> evalFlag(
108+
flag: Flag<T>,
109+
context: Context,
110+
): String? {
111+
return httpClient.newCall(
72112
Request.Builder()
73-
.url("$baseurl/${flag.key}")
113+
.url("$urlPrefix/${flag.key}/evaluate?kind=${context.kind}&value=${context.key}")
74114
.build(),
75-
).execute()
115+
).execute().body?.string()
76116
}
77117

78-
fun <T> evalFlag(
118+
fun <T> setRule(
119+
flag: Flag<T>,
120+
context: Context,
121+
value: T,
122+
) {
123+
val requestRule =
124+
ApiRule(
125+
context = ApiContext(kind = context.kind, value = context.key),
126+
value = value.toString(),
127+
)
128+
val response =
129+
httpClient.newCall(
130+
Request.Builder()
131+
.url("$urlPrefix/${flag.key}/rules")
132+
.post(Jsons.serialize(requestRule).toRequestBody("application/json".toMediaType()))
133+
.build(),
134+
).execute()
135+
assert(response.code == 200, { "Failed to update the feature flag ${flag.key}, error: ${response.code}: ${response.body?.string()}" })
136+
}
137+
138+
fun <T> deleteRule(
79139
flag: Flag<T>,
80140
context: Context,
81141
) {
142+
val requestContext = ApiContext(kind = context.kind, value = context.key)
82143
httpClient.newCall(
83144
Request.Builder()
84-
.url("$baseurl/${flag.key}/evaluate?kind=${context.kind}&value=${context.key}")
145+
.url("$urlPrefix/${flag.key}/rules")
146+
.delete(Jsons.serialize(requestContext).toRequestBody("application/json".toMediaType()))
85147
.build(),
86148
).execute()
87149
}

‎airbyte-test-utils/build.gradle.kts

+1
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ dependencies {
2525
implementation(project(":oss:airbyte-commons-storage"))
2626
implementation(project(":oss:airbyte-commons-temporal"))
2727
implementation(project(":oss:airbyte-commons-worker"))
28+
implementation(project(":oss:airbyte-featureflag"))
2829
implementation(libs.bundles.kubernetes.client)
2930
implementation(libs.bundles.flyway)
3031
implementation(libs.temporal.sdk)

‎airbyte-test-utils/src/main/java/io/airbyte/test/utils/AcceptanceTestHarness.java

+24
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@
100100
import io.airbyte.db.Database;
101101
import io.airbyte.db.factory.DataSourceFactory;
102102
import io.airbyte.db.jdbc.JdbcUtils;
103+
import io.airbyte.featureflag.tests.TestFlagsSetter;
103104
import io.airbyte.test.container.AirbyteTestContainer;
104105
import io.temporal.client.WorkflowClient;
105106
import io.temporal.serviceclient.WorkflowServiceStubs;
@@ -226,6 +227,7 @@ public class AcceptanceTestHarness {
226227

227228
private AirbyteTestContainer airbyteTestContainer;
228229
private final AirbyteApiClient apiClient;
230+
private final TestFlagsSetter testFlagsSetter;
229231
private final UUID defaultWorkspaceId;
230232
private final String postgresSqlInitFile;
231233

@@ -253,10 +255,19 @@ public void removeConnection(final UUID connection) {
253255
public AcceptanceTestHarness(final AirbyteApiClient apiClient,
254256
final UUID defaultWorkspaceId,
255257
final String postgresSqlInitFile)
258+
throws GeneralSecurityException, URISyntaxException, IOException, InterruptedException {
259+
this(apiClient, null, defaultWorkspaceId, postgresSqlInitFile);
260+
}
261+
262+
public AcceptanceTestHarness(final AirbyteApiClient apiClient,
263+
final TestFlagsSetter testFlagsSetter,
264+
final UUID defaultWorkspaceId,
265+
final String postgresSqlInitFile)
256266
throws URISyntaxException, IOException, InterruptedException, GeneralSecurityException {
257267
// reads env vars to assign static variables
258268
assignEnvVars();
259269
this.apiClient = apiClient;
270+
this.testFlagsSetter = testFlagsSetter;
260271
this.defaultWorkspaceId = defaultWorkspaceId;
261272
this.postgresSqlInitFile = postgresSqlInitFile;
262273

@@ -325,6 +336,19 @@ public AcceptanceTestHarness(final AirbyteApiClient apiClient, final UUID defaul
325336
this(apiClient, defaultWorkspaceId, DEFAULT_POSTGRES_INIT_SQL_FILE);
326337
}
327338

339+
public AcceptanceTestHarness(final AirbyteApiClient apiClient, final UUID defaultWorkspaceId, final TestFlagsSetter testFlagsSetter)
340+
throws GeneralSecurityException, URISyntaxException, IOException, InterruptedException {
341+
this(apiClient, testFlagsSetter, defaultWorkspaceId, DEFAULT_POSTGRES_INIT_SQL_FILE);
342+
}
343+
344+
public AirbyteApiClient getApiClient() {
345+
return apiClient;
346+
}
347+
348+
public TestFlagsSetter getTestFlagsSetter() {
349+
return testFlagsSetter;
350+
}
351+
328352
public void stopDbAndContainers() {
329353
if (isGke) {
330354
try {

‎airbyte-tests/src/test-acceptance/java/io/airbyte/test/acceptance/AcceptanceTestsResources.java

+9-3
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import io.airbyte.api.client.model.generated.WorkspaceCreate;
3636
import io.airbyte.commons.json.Jsons;
3737
import io.airbyte.db.Database;
38+
import io.airbyte.featureflag.tests.TestFlagsSetter;
3839
import io.airbyte.test.utils.AcceptanceTestHarness;
3940
import io.airbyte.test.utils.Asserts;
4041
import io.airbyte.test.utils.TestConnectionCreate.Builder;
@@ -240,8 +241,10 @@ void runIncrementalSyncForAWorkspaceId(final UUID workspaceId) throws Exception
240241
StreamStatusJobType.SYNC);
241242
}
242243

243-
void runSmallSyncForAWorkspaceId(final UUID workspaceId) throws Exception {
244-
LOGGER.info("Starting runSmallSyncForAWorkspaceId()");
244+
record SyncIds(UUID connectionId, Long jobId, Integer attemptNumber) {}
245+
246+
SyncIds runSmallSyncForAWorkspaceId(final UUID workspaceId) throws Exception {
247+
LOGGER.info("Starting runSmallSyncForAWorkspaceId(" + workspaceId + ")");
245248
final UUID sourceId = testHarness.createPostgresSource(workspaceId).getSourceId();
246249
final UUID destinationId = testHarness.createPostgresDestination(workspaceId).getDestinationId();
247250
final SourceDiscoverSchemaRead discoverResult = testHarness.discoverSourceSchemaWithId(sourceId);
@@ -289,12 +292,15 @@ void runSmallSyncForAWorkspaceId(final UUID workspaceId) throws Exception {
289292
WITHOUT_SCD_TABLE);
290293
Asserts.assertStreamStatuses(testHarness, workspaceId, connectionId, connectionSyncRead1.getJob().getId(), StreamStatusRunState.COMPLETE,
291294
StreamStatusJobType.SYNC);
295+
296+
return new SyncIds(connectionId, connectionSyncRead1.getJob().getId(), connectionSyncRead1.getAttempts().size() - 1);
292297
}
293298

294299
void init() throws URISyntaxException, IOException, InterruptedException, GeneralSecurityException {
295300
final AirbyteApiClient airbyteApiClient =
296301
createAirbyteApiClient(AIRBYTE_SERVER_HOST + "/api",
297302
Map.of(GATEWAY_AUTH_HEADER, CLOUD_API_USER_HEADER_VALUE));
303+
final TestFlagsSetter testFlagsSetter = new TestFlagsSetter(AIRBYTE_SERVER_HOST);
298304

299305
// If a workspace id is passed, use that. Otherwise, create a new workspace.
300306
// NOTE: we want to sometimes use a pre-configured workspace e.g., if we run against a production
@@ -326,7 +332,7 @@ void init() throws URISyntaxException, IOException, InterruptedException, Genera
326332
LOGGER.info("pg source definition: {}", sourceDef.getDockerImageTag());
327333
LOGGER.info("pg destination definition: {}", destinationDef.getDockerImageTag());
328334

329-
testHarness = new AcceptanceTestHarness(airbyteApiClient, workspaceId);
335+
testHarness = new AcceptanceTestHarness(airbyteApiClient, workspaceId, testFlagsSetter);
330336

331337
testHarness.ensureCleanSlate();
332338
}

0 commit comments

Comments
 (0)