From 351d5215268a7eec2b63f52da643628820f19f4f Mon Sep 17 00:00:00 2001
From: Tobias Happ <tobias.happ@gmx.de>
Date: Fri, 30 Apr 2021 21:30:43 +0200
Subject: [PATCH] Add flushEnabled option for couchbase buckets

---
 .../couchbase/BucketDefinition.java           | 16 +++++
 .../couchbase/CouchbaseContainer.java         |  1 +
 .../couchbase/CouchbaseContainerTest.java     | 58 ++++++++++++++-----
 3 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/modules/couchbase/src/main/java/org/testcontainers/couchbase/BucketDefinition.java b/modules/couchbase/src/main/java/org/testcontainers/couchbase/BucketDefinition.java
index 8c2b11403ad..3778833182e 100644
--- a/modules/couchbase/src/main/java/org/testcontainers/couchbase/BucketDefinition.java
+++ b/modules/couchbase/src/main/java/org/testcontainers/couchbase/BucketDefinition.java
@@ -22,6 +22,7 @@
 public class BucketDefinition {
 
     private final String name;
+    private boolean flushEnabled = false;
     private boolean queryPrimaryIndex = true;
     private int quota = 100;
 
@@ -29,6 +30,17 @@ public BucketDefinition(final String name) {
         this.name = name;
     }
 
+    /**
+     * Enables flush for this bucket (disabled by default).
+     *
+     * @param flushEnabled if true, the bucket can be flushed.
+     * @return this {@link BucketDefinition} for chaining purposes.
+     */
+    public BucketDefinition withFlushEnabled(final boolean flushEnabled) {
+        this.flushEnabled = flushEnabled;
+        return this;
+    }
+
     /**
      * Sets a custom bucket quota (100MB by default).
      *
@@ -58,6 +70,10 @@ public String getName() {
         return name;
     }
 
+    public boolean hasFlushEnabled() {
+        return flushEnabled;
+    }
+
     public boolean hasPrimaryIndex() {
         return queryPrimaryIndex;
     }
diff --git a/modules/couchbase/src/main/java/org/testcontainers/couchbase/CouchbaseContainer.java b/modules/couchbase/src/main/java/org/testcontainers/couchbase/CouchbaseContainer.java
index 8df5844f538..64489f30531 100644
--- a/modules/couchbase/src/main/java/org/testcontainers/couchbase/CouchbaseContainer.java
+++ b/modules/couchbase/src/main/java/org/testcontainers/couchbase/CouchbaseContainer.java
@@ -368,6 +368,7 @@ private void createBuckets() {
             @Cleanup Response response = doHttpRequest(MGMT_PORT, "/pools/default/buckets", "POST", new FormBody.Builder()
                 .add("name", bucket.getName())
                 .add("ramQuotaMB", Integer.toString(bucket.getQuota()))
+                .add("flushEnabled", bucket.hasFlushEnabled() ? "1" : "0")
                 .build(), true);
 
             checkSuccessfulResponse(response, "Could not create bucket " + bucket.getName());
diff --git a/modules/couchbase/src/test/java/org/testcontainers/couchbase/CouchbaseContainerTest.java b/modules/couchbase/src/test/java/org/testcontainers/couchbase/CouchbaseContainerTest.java
index 13709299372..6b56e9079d1 100644
--- a/modules/couchbase/src/test/java/org/testcontainers/couchbase/CouchbaseContainerTest.java
+++ b/modules/couchbase/src/test/java/org/testcontainers/couchbase/CouchbaseContainerTest.java
@@ -23,7 +23,10 @@
 import org.junit.Test;
 import org.testcontainers.utility.DockerImageName;
 
+import java.util.function.Consumer;
+
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 
 public class CouchbaseContainerTest {
 
@@ -41,17 +44,7 @@ public void testBasicContainerUsage() {
                 .withBucket(bucketDefinition)
             // }
         ) {
-            container.start();
-
-            // cluster_creation {
-            Cluster cluster = Cluster.connect(
-                container.getConnectionString(),
-                container.getUsername(),
-                container.getPassword()
-            );
-            // }
-
-            try {
+            setUpClient(container, cluster -> {
                 Bucket bucket = cluster.bucket(bucketDefinition.getName());
                 Collection collection = bucket.defaultCollection();
 
@@ -60,10 +53,47 @@ public void testBasicContainerUsage() {
                 JsonObject fooObject = collection.get("foo").contentAsObject();
 
                 assertEquals("value", fooObject.getString("key"));
-            } finally {
-                cluster.disconnect();
-            }
+            });
         }
     }
 
+    @Test
+    public void testBucketIsFlushableIfEnabled() {
+        BucketDefinition bucketDefinition = new BucketDefinition("mybucket")
+            .withFlushEnabled(true);
+
+        try (
+            CouchbaseContainer container = new CouchbaseContainer(COUCHBASE_IMAGE)
+                .withBucket(bucketDefinition)
+        ) {
+            setUpClient(container, cluster -> {
+                Bucket bucket = cluster.bucket(bucketDefinition.getName());
+                Collection collection = bucket.defaultCollection();
+
+                collection.upsert("foo", JsonObject.create().put("key", "value"));
+
+                cluster.buckets().flushBucket(bucketDefinition.getName());
+
+                assertFalse(collection.exists("foo").exists());
+            });
+        }
+    }
+
+    private void setUpClient(CouchbaseContainer container, Consumer<Cluster> consumer) {
+        container.start();
+
+        // cluster_creation {
+        Cluster cluster = Cluster.connect(
+            container.getConnectionString(),
+            container.getUsername(),
+            container.getPassword()
+        );
+        // }
+
+        try {
+            consumer.accept(cluster);
+        } finally {
+            cluster.disconnect();
+        }
+    }
 }