Skip to content

Commit

Permalink
Changed wait strategy for ysql testcontainer (#1)
Browse files Browse the repository at this point in the history 
* Changed wait strategy for ysql testcontainer

* Updating YugabyteDB-YSQL WaitStrategy

* Catch IllegalStateException

* Update test query

* Changes as per review

* Added a test case.

* Updating YSQL wait strategy

* Removing redundant if.

* Minor Change

* YugabyteDB version updates in tests
  • Loading branch information
@HarshDaryani896
HarshDaryani896 authored Jul 27, 2023
1 parent 8465795 commit a70f002
Show file tree
Hide file tree
Showing 6 changed files with 48 additions and 18 deletions.
18 changes: 17 additions & 1 deletion modules/yugabytedb/src/main/java/org/testcontainers/containers/YugabyteDBYSQLContainer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ public class YugabyteDBYSQLContainer extends JdbcDatabaseContainer<YugabyteDBYSQ

private static final String ENTRYPOINT = "bin/yugabyted start --background=false";

private boolean extendedStartupProbe = true;

private String database = "yugabyte";

private String username = "yugabyte";
Expand All @@ -49,7 +51,7 @@ public YugabyteDBYSQLContainer(final DockerImageName imageName) {
super(imageName);
imageName.assertCompatibleWith(DEFAULT_IMAGE_NAME);
withExposedPorts(YSQL_PORT, MASTER_DASHBOARD_PORT, TSERVER_DASHBOARD_PORT);
waitingFor(new YugabyteDBYSQLWaitStrategy(this).withStartupTimeout(Duration.ofSeconds(60)));
waitingFor(new YugabyteDBYSQLWaitStrategy().withStartupTimeout(Duration.ofSeconds(60)));
withCommand(ENTRYPOINT);
}

Expand Down Expand Up @@ -96,6 +98,10 @@ public String getDatabaseName() {
return database;
}

public boolean isExtendedStartupProbe() {
return extendedStartupProbe;
}

@Override
public String getUsername() {
return username;
Expand Down Expand Up @@ -123,6 +129,11 @@ public YugabyteDBYSQLContainer withDatabaseName(final String database) {
return this;
}

public YugabyteDBYSQLContainer withExtendedStartupProbe(final boolean startupProbe) {
this.extendedStartupProbe = startupProbe;
return this;
}

/**
* Setting this would create the custom user role
* @param username user name
Expand All @@ -146,4 +157,9 @@ public YugabyteDBYSQLContainer withPassword(final String password) {
this.password = password;
return this;
}

@Override
protected void waitUntilContainerStarted() {
getWaitStrategy().waitUntilReady(this);
}
}
2 changes: 2 additions & 0 deletions ...gabytedb/src/main/java/org/testcontainers/containers/delegate/YugabyteDBYCQLDelegate.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,11 @@ public void execute(
boolean continueOnError,
boolean ignoreFailedDrops
) {
final String containerInterfaceIP = container.getContainerInfo().getNetworkSettings().getNetworks().entrySet().stream().findFirst().get().getValue().getIpAddress();
try {
ExecResult result = container.execInContainer(
BIN_PATH,
containerInterfaceIP,
"-u",
container.getUsername(),
"-p",
Expand Down
2 changes: 2 additions & 0 deletions ...tedb/src/main/java/org/testcontainers/containers/strategy/YugabyteDBYCQLWaitStrategy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ public final class YugabyteDBYCQLWaitStrategy extends AbstractWaitStrategy {
public void waitUntilReady(WaitStrategyTarget target) {
YugabyteDBYCQLContainer container = (YugabyteDBYCQLContainer) target;
AtomicBoolean status = new AtomicBoolean(true);
final String containerInterfaceIP = container.getContainerInfo().getNetworkSettings().getNetworks().entrySet().stream().findFirst().get().getValue().getIpAddress();
retryUntilSuccess(
(int) startupTimeout.getSeconds(),
TimeUnit.SECONDS,
Expand All @@ -46,6 +47,7 @@ public void waitUntilReady(WaitStrategyTarget target) {
try {
ExecResult result = container.execInContainer(
BIN_PATH,
containerInterfaceIP,
"-u",
container.getUsername(),
"-p",
Expand Down
22 changes: 9 additions & 13 deletions ...tedb/src/main/java/org/testcontainers/containers/strategy/YugabyteDBYSQLWaitStrategy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@
import lombok.extern.slf4j.Slf4j;
import org.testcontainers.containers.YugabyteDBYSQLContainer;
import org.testcontainers.containers.wait.strategy.AbstractWaitStrategy;
import org.testcontainers.containers.wait.strategy.WaitStrategyTarget;

import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.concurrent.TimeUnit;

import static org.rnorth.ducttape.unreliables.Unreliables.retryUntilSuccess;
Expand All @@ -29,30 +29,26 @@ public final class YugabyteDBYSQLWaitStrategy extends AbstractWaitStrategy {

private static final String YSQL_TEST_QUERY = "SELECT 1";

private final WaitStrategyTarget target;
private static final String YSQL_EXTENDED_PROBE = "CREATE TEMP TABLE IF NOT EXISTS YB_SAMPLE(k int, v int, primary key(k, v))";

@Override
public void waitUntilReady(WaitStrategyTarget target) {
YugabyteDBYSQLContainer container = (YugabyteDBYSQLContainer) target;
public void waitUntilReady() {
YugabyteDBYSQLContainer container = (YugabyteDBYSQLContainer) waitStrategyTarget;
retryUntilSuccess(
(int) startupTimeout.getSeconds(),
TimeUnit.SECONDS,
() -> {
getRateLimiter()
.doWhenReady(() -> {
try (Connection con = container.createConnection(container.getJdbcUrl())) {
con.createStatement().execute(YSQL_TEST_QUERY);
} catch (SQLException ex) {
log.error("Error connecting to the database", ex);
try (Connection con = container.createConnection(""); Statement stmt = con.createStatement()) {
stmt.execute(container.isExtendedStartupProbe() ? YSQL_EXTENDED_PROBE : YSQL_TEST_QUERY);
}
catch (SQLException ex) {
throw new RuntimeException(ex);
}
});
return true;
}
);
}

@Override
public void waitUntilReady() {
waitUntilReady(target);
}
}
4 changes: 2 additions & 2 deletions modules/yugabytedb/src/test/java/org/testcontainers/junit/yugabytedb/YugabyteDBYCQLTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
*/
public class YugabyteDBYCQLTest {

private static final String IMAGE_NAME = "yugabytedb/yugabyte:2.14.4.0-b26";
private static final String IMAGE_NAME = "yugabytedb/yugabyte:2.18.1.0-b84";

private static final DockerImageName YBDB_TEST_IMAGE = DockerImageName.parse(IMAGE_NAME);

Expand All @@ -22,7 +22,7 @@ public void testSmoke() {
try (
// creatingYCQLContainer {
final YugabyteDBYCQLContainer ycqlContainer = new YugabyteDBYCQLContainer(
"yugabytedb/yugabyte:2.14.4.0-b26"
"yugabytedb/yugabyte:2.18.1.0-b84"
)
// }
) {
Expand Down
18 changes: 16 additions & 2 deletions modules/yugabytedb/src/test/java/org/testcontainers/junit/yugabytedb/YugabyteDBYSQLTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
*/
public class YugabyteDBYSQLTest extends AbstractContainerDatabaseTest {

private static final String IMAGE_NAME = "yugabytedb/yugabyte:2.14.4.0-b26";
private static final String IMAGE_NAME = "yugabytedb/yugabyte:2.18.1.0-b84";

private static final DockerImageName YBDB_TEST_IMAGE = DockerImageName.parse(IMAGE_NAME);

Expand All @@ -23,7 +23,7 @@ public void testSmoke() throws SQLException {
try (
// creatingYSQLContainer {
final YugabyteDBYSQLContainer ysqlContainer = new YugabyteDBYSQLContainer(
"yugabytedb/yugabyte:2.14.4.0-b26"
"yugabytedb/yugabyte:2.18.1.0-b84"
)
// }
) {
Expand Down Expand Up @@ -95,4 +95,18 @@ public void testWithCustomRole() throws SQLException {
.isEqualTo(1);
}
}

@Test
public void testExtendedProbe() throws SQLException {
try (
final YugabyteDBYSQLContainer ysqlContainer = new YugabyteDBYSQLContainer(YBDB_TEST_IMAGE)
.withExtendedStartupProbe(false)
.withInitScript("init/init_yql.sql")
) {
ysqlContainer.start();
assertThat(performQuery(ysqlContainer, "SELECT greet FROM dsql").getString(1))
.as("A record match succeeds")
.isEqualTo("Hello DSQL");
}
}
}

0 comments on commit a70f002

Please sign in to comment.