From 96cf316791779cf00b7e5bee468102e284e81720 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Wed, 12 Mar 2025 08:37:42 -0700
Subject: [PATCH] Internal Changes

PiperOrigin-RevId: 736144368
---
 .../main/java/dev/cel/bundle/CelEnvironment.java | 16 ++++++++++++++--
 .../dev/cel/bundle/CelEnvironmentYamlParser.java |  2 +-
 .../java/dev/cel/bundle/CelEnvironmentTest.java  |  2 +-
 .../cel/bundle/CelEnvironmentYamlParserTest.java |  6 +++---
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/bundle/src/main/java/dev/cel/bundle/CelEnvironment.java b/bundle/src/main/java/dev/cel/bundle/CelEnvironment.java
index 9010a5307..1c1c05ee1 100644
--- a/bundle/src/main/java/dev/cel/bundle/CelEnvironment.java
+++ b/bundle/src/main/java/dev/cel/bundle/CelEnvironment.java
@@ -101,6 +101,8 @@ public abstract class CelEnvironment {
   @AutoValue.Builder
   public abstract static class Builder {
 
+    public abstract ImmutableSet.Builder<ExtensionConfig> extensionsBuilder();
+
     // For testing only, to empty out the source.
     abstract Builder setSource(Optional<Source> source);
 
@@ -112,7 +114,18 @@ public abstract static class Builder {
 
     public abstract Builder setContainer(String container);
 
-    public abstract Builder setExtensions(ImmutableSet<ExtensionConfig> extensions);
+    @CanIgnoreReturnValue
+    public Builder addExtensions(ExtensionConfig... extensions) {
+      checkNotNull(extensions);
+      return addExtensions(Arrays.asList(extensions));
+    }
+
+    @CanIgnoreReturnValue
+    public Builder addExtensions(Iterable<ExtensionConfig> extensions) {
+      checkNotNull(extensions);
+      this.extensionsBuilder().addAll(extensions);
+      return this;
+    }
 
     @CanIgnoreReturnValue
     public Builder setVariables(VariableDecl... variables) {
@@ -138,7 +151,6 @@ public static Builder newBuilder() {
         .setName("")
         .setDescription("")
         .setContainer("")
-        .setExtensions(ImmutableSet.of())
         .setVariables(ImmutableSet.of())
         .setFunctions(ImmutableSet.of());
   }
diff --git a/bundle/src/main/java/dev/cel/bundle/CelEnvironmentYamlParser.java b/bundle/src/main/java/dev/cel/bundle/CelEnvironmentYamlParser.java
index ab9a9ff54..64cdda0bf 100644
--- a/bundle/src/main/java/dev/cel/bundle/CelEnvironmentYamlParser.java
+++ b/bundle/src/main/java/dev/cel/bundle/CelEnvironmentYamlParser.java
@@ -475,7 +475,7 @@ private CelEnvironment.Builder parseConfig(ParserContext<Node> ctx, Node node) {
             builder.setFunctions(parseFunctions(ctx, valueNode));
             break;
           case "extensions":
-            builder.setExtensions(parseExtensions(ctx, valueNode));
+            builder.addExtensions(parseExtensions(ctx, valueNode));
             break;
           default:
             ctx.reportError(id, "Unknown config tag: " + fieldName);
diff --git a/bundle/src/test/java/dev/cel/bundle/CelEnvironmentTest.java b/bundle/src/test/java/dev/cel/bundle/CelEnvironmentTest.java
index 3f5a69e9f..40ee252d0 100644
--- a/bundle/src/test/java/dev/cel/bundle/CelEnvironmentTest.java
+++ b/bundle/src/test/java/dev/cel/bundle/CelEnvironmentTest.java
@@ -54,7 +54,7 @@ public void extend_allExtensions() throws Exception {
             ExtensionConfig.of("sets"),
             ExtensionConfig.of("strings"));
     CelEnvironment environment =
-        CelEnvironment.newBuilder().setExtensions(extensionConfigs).build();
+        CelEnvironment.newBuilder().addExtensions(extensionConfigs).build();
 
     Cel cel = environment.extend(CelFactory.standardCelBuilder().build(), CelOptions.DEFAULT);
     CelAbstractSyntaxTree ast =
diff --git a/bundle/src/test/java/dev/cel/bundle/CelEnvironmentYamlParserTest.java b/bundle/src/test/java/dev/cel/bundle/CelEnvironmentYamlParserTest.java
index 34e1e2f6b..90e6b47c3 100644
--- a/bundle/src/test/java/dev/cel/bundle/CelEnvironmentYamlParserTest.java
+++ b/bundle/src/test/java/dev/cel/bundle/CelEnvironmentYamlParserTest.java
@@ -98,7 +98,7 @@ public void environment_setExtensions() throws Exception {
         .isEqualTo(
             CelEnvironment.newBuilder()
                 .setSource(environment.source().get())
-                .setExtensions(
+                .addExtensions(
                     ImmutableSet.of(
                         ExtensionConfig.of("bindings"),
                         ExtensionConfig.of("encoders"),
@@ -125,7 +125,7 @@ public void environment_setExtensionVersionToLatest() throws Exception {
         .isEqualTo(
             CelEnvironment.newBuilder()
                 .setSource(environment.source().get())
-                .setExtensions(ImmutableSet.of(ExtensionConfig.of("bindings", Integer.MAX_VALUE)))
+                .addExtensions(ImmutableSet.of(ExtensionConfig.of("bindings", Integer.MAX_VALUE)))
                 .build());
     assertThat(environment.extend(CEL_WITH_MESSAGE_TYPES, CelOptions.DEFAULT)).isNotNull();
   }
@@ -675,7 +675,7 @@ private enum EnvironmentYamlResourceTestCase {
         CelEnvironment.newBuilder()
             .setName("extended-env")
             .setContainer("cel.expr")
-            .setExtensions(
+            .addExtensions(
                 ImmutableSet.of(
                     ExtensionConfig.of("optional", 2),
                     ExtensionConfig.of("math", Integer.MAX_VALUE)))