Add test helpers to create and delete API keys (#2193)

- Add test helpers to create and delete API keys - Allow `APIContainerHelper.deleteContainer` to handle 504 response - Add `waitFor` methods to test `Timer`
LabKey · Dec 17, 2024 · 7f9f30d · 7f9f30d
1 parent e07db68
commit 7f9f30d
Show file tree

Hide file tree

Showing 6 changed files with 153 additions and 37 deletions.
diff --git a/src/org/labkey/test/WebDriverWrapper.java b/src/org/labkey/test/WebDriverWrapper.java
@@ -49,12 +49,12 @@
 import org.labkey.test.selenium.EphemeralWebElement;
 import org.labkey.test.util.CodeMirrorHelper;
 import org.labkey.test.util.Crawler;
-import org.labkey.test.util.OptionalFeatureHelper;
 import org.labkey.test.util.Ext4Helper;
 import org.labkey.test.util.ExtHelper;
 import org.labkey.test.util.LabKeyExpectedConditions;
 import org.labkey.test.util.LogMethod;
 import org.labkey.test.util.LoggedParam;
+import org.labkey.test.util.OptionalFeatureHelper;
 import org.labkey.test.util.PasswordUtil;
 import org.labkey.test.util.RelativeUrl;
 import org.labkey.test.util.TestLogger;
@@ -2265,21 +2265,12 @@ public WebElement doAndWaitForElementToRefresh(Runnable func, Locator loc, WebDr
      * Wait for Supplier to return non-null non-false value
      * @param wait milliseconds
      * @return final result of Supplier.get()
+     * @see Timer#waitFor(Supplier)
      */
     @Contract(pure = true)
     public static <T> T waitFor(Supplier<T> checker, int wait)
     {
-        long startTime = System.currentTimeMillis();
-        T result;
-        do
-        {
-            result = checker.get();
-            if (result != null && !Boolean.FALSE.equals(result))
-                break;
-            sleep(100);
-        } while ((System.currentTimeMillis() - startTime) < wait);
-
-        return result;
+        return new Timer(Duration.ofMillis(wait)).waitFor(checker);
     }
 
     public static void waitForEquals(String message, Supplier<?> expected, Supplier<?> actual, int wait)
@@ -2298,17 +2289,20 @@ public static void waitForNotEquals(String message, Supplier<?> expected, Suppli
         }
     }
 
-    public static void waitFor(Supplier<Boolean> checker, String failMessage, int wait)
+    /**
+     * @see Timer#waitFor(Supplier, String)
+     */
+    public static <T> T waitFor(Supplier<T> checker, String failMessage, int wait)
     {
-        waitFor(checker, () -> failMessage, wait);
+        return new Timer(Duration.ofMillis(wait)).waitFor(checker, failMessage);
     }
 
-    public static void waitFor(Supplier<Boolean> checker, Supplier<String> failMessage, int wait)
+    /**
+     * @see Timer#waitFor(Supplier, Supplier)
+     */
+    public static <T> T waitFor(Supplier<T> checker, Supplier<String> failMessageSupplier, int wait)
     {
-        if (!waitFor(checker, wait))
-        {
-            throw new TimeoutException(failMessage.get() + TestLogger.formatElapsedTime(wait));
-        }
+        return new Timer(Duration.ofMillis(wait)).waitFor(checker, failMessageSupplier);
     }
 
     public File clickAndWaitForDownload(Locator elementToClick)

diff --git a/src/org/labkey/test/WebTestHelper.java b/src/org/labkey/test/WebTestHelper.java
@@ -46,6 +46,9 @@
 import org.labkey.remoteapi.CommandResponse;
 import org.labkey.remoteapi.Connection;
 import org.labkey.remoteapi.SimplePostCommand;
+import org.labkey.remoteapi.query.DeleteRowsCommand;
+import org.labkey.remoteapi.query.Filter;
+import org.labkey.remoteapi.query.SelectRowsCommand;
 import org.labkey.serverapi.reader.Readers;
 import org.labkey.test.util.InstallCert;
 import org.labkey.test.util.LogMethod;
@@ -79,9 +82,11 @@
 import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Random;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -107,6 +112,7 @@ public class WebTestHelper
     private static boolean USE_CONTAINER_RELATIVE_URL = true;
     private static final Map<String, Map<String, Cookie>> savedCookies = new HashMap<>();
     private static final Map<String, String> savedSessionKeys = new HashMap<>();
+    private static final Map<String, String> savedApiKeys = new HashMap<>();
 
     static { TestProperties.load(); }
 
@@ -139,31 +145,73 @@ public static String getSessionKey(String user)
         }
         String sessionId = sessionCookie.getValue();
 
-        if (!savedSessionKeys.containsKey(sessionId))
+        return savedSessionKeys.computeIfAbsent(sessionId, k -> createApiKey(getRemoteApiConnection(user, true), "session", null));
+    }
+
+    public static String createApiKey(Connection connection)
+    {
+        String description = UUID.randomUUID().toString(); // Something unique that we can query for deletion
+        String apiKey = createApiKey(connection, API_KEY, description);
+        savedApiKeys.put(apiKey, description);
+        return apiKey;
+    }
+
+    private static String createApiKey(Connection connection, String type, String description)
+    {
+        SimplePostCommand command = new SimplePostCommand("security", "createApiKey");
+        JSONObject json = new JSONObject();
+        json.put("type", type);
+        json.put("description", description);
+        command.setJsonObject(json);
+
+        try
+        {
+            CommandResponse response = command.execute(connection, "/");
+            String apikey = (String) response.getParsedData().get("apikey");
+            if (apikey == null)
+            {
+                TestLogger.error(response.getText());
+                throw new RuntimeException("Failed to generate %s key".formatted(type));
+            }
+            return apikey;
+        }
+        catch (CommandException | IOException e)
         {
-            Connection connection = getRemoteApiConnection(user, true);
-            SimplePostCommand command = new SimplePostCommand("security", "createApiKey");
-            JSONObject json = new JSONObject();
-            json.put("type", "session");
-            command.setJsonObject(json);
+            throw new RuntimeException("Failed to generate %s key".formatted(type), e);
+        }
+    }
 
+    public static void deleteApiKey(Connection connection, String apiKey)
+    {
+        String description = savedApiKeys.get(apiKey);
+        if (description != null)
+        {
             try
             {
-                CommandResponse response = command.execute(connection, "/");
-                Object apikey = response.getParsedData().get("apikey");
-                if (apikey == null)
+                SelectRowsCommand selectRowsCommand = new SelectRowsCommand("core", "apiKeys");
+                selectRowsCommand.setFilters(List.of(new Filter("Description", description)));
+                List<Map<String, Object>> rows = selectRowsCommand.execute(connection, null).getRows();
+                if (rows.size() == 1)
                 {
-                    TestLogger.error(response.getText());
-                    throw new RuntimeException("Failed to generate session key");
+                    DeleteRowsCommand deleteRowsCommand = new DeleteRowsCommand("core", "apiKeys");
+                    deleteRowsCommand.setRows(rows);
+                    deleteRowsCommand.execute(connection, null);
+                    savedApiKeys.remove(apiKey);
+                }
+                else
+                {
+                    TestLogger.log("Skipping apiKey deletion. Unexpected number of rows found: " + rows.size());
                 }
-                savedSessionKeys.put(sessionId, (String) apikey);
             }
-            catch (CommandException | IOException e)
+            catch (IOException | CommandException e)
             {
-                throw new RuntimeException("Unable to generate session key", e);
+                throw new RuntimeException(e);
             }
         }
-        return savedSessionKeys.get(sessionId);
+        else
+        {
+            TestLogger.warn("Refusing to delete an API key not created by this test");
+        }
     }
 
     public static boolean isUseContainerRelativeUrl()

diff --git a/src/org/labkey/test/stress/AbstractScenario.java b/src/org/labkey/test/stress/AbstractScenario.java
@@ -4,6 +4,8 @@
 import org.jetbrains.annotations.Nullable;
 import org.junit.Assert;
 import org.labkey.remoteapi.Connection;
+import org.labkey.test.TestProperties;
+import org.labkey.test.teamcity.TeamCityUtils;
 import org.labkey.test.util.TestDateUtils;
 import org.labkey.test.util.TestLogger;
 
@@ -183,6 +185,20 @@ public final Set<T> finishScenario() throws InterruptedException
             shutdownNow();
             throw t;
         }
+        finally
+        {
+            if (getResultsFile() != null)
+            {
+                if (TestProperties.isTestRunningOnTeamCity())
+                {
+                    TeamCityUtils.publishArtifact(getResultsFile(), null);
+                }
+                else
+                {
+                    TestLogger.log("Perf data available in: " + getResultsFile().getAbsolutePath());
+                }
+            }
+        }
     }
 
     /**

diff --git a/src/org/labkey/test/util/APIContainerHelper.java b/src/org/labkey/test/util/APIContainerHelper.java
@@ -27,13 +27,15 @@
 import org.labkey.remoteapi.security.DeleteContainerCommand;
 import org.labkey.remoteapi.security.GetContainersCommand;
 import org.labkey.test.BaseWebDriverTest;
+import org.labkey.test.TestProperties;
 import org.labkey.test.TestTimeoutException;
 import org.labkey.test.WebTestHelper;
 
 import java.io.IOException;
 import java.net.SocketTimeoutException;
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -147,15 +149,14 @@ public void deleteContainer(String path, boolean failIfNotFound, int wait) throw
     {
         WebTestHelper.logToServer("=Test= Starting container delete: " + path);
 
+        Timer deleteTimer = new Timer(Duration.ofMillis(wait));
         DeleteContainerCommand dcc = new DeleteContainerCommand();
         dcc.setTimeout(wait);
         try
         {
             Connection defaultConnection = _test.createDefaultConnection();
             defaultConnection.setTimeout(wait);
             dcc.execute(defaultConnection, path);
-
-            WebTestHelper.logToServer("=Test= Finished container delete: " + path);
         }
         catch (CommandException e)
         {
@@ -164,6 +165,14 @@ public void deleteContainer(String path, boolean failIfNotFound, int wait) throw
                 if (failIfNotFound)
                     fail("Container not found: " + path);
             }
+            else if (TestProperties.isServerRemote() && e.getStatusCode() == HttpStatus.SC_GATEWAY_TIMEOUT)
+            {
+                TestLogger.log("Waiting for container deletion after Gateway Timeout (504) error");
+                deleteTimer.waitFor(() -> !doesContainerExist(path), () -> {
+                    WebTestHelper.logToServer("=Test= Timed out deleting container: " + path);
+                    return "Timed out deleting container [%s]".formatted(path);
+                });
+            }
             else
             {
                 throw new RuntimeException("Failed to delete container: " + path, e);
@@ -178,6 +187,8 @@ public void deleteContainer(String path, boolean failIfNotFound, int wait) throw
         {
             throw new RuntimeException("Failed to delete container: " + path, e);
         }
+
+        WebTestHelper.logToServer("=Test= Finished container delete: " + path);
     }
 
     @Override

diff --git a/src/org/labkey/test/util/Timer.java b/src/org/labkey/test/util/Timer.java
@@ -16,11 +16,16 @@
 package org.labkey.test.util;
 
 import org.apache.commons.lang3.time.StopWatch;
+import org.jetbrains.annotations.Contract;
+import org.openqa.selenium.TimeoutException;
 
 import java.time.Duration;
 import java.time.Instant;
 import java.time.LocalDateTime;
 import java.time.ZoneId;
+import java.util.function.Supplier;
+
+import static org.labkey.test.WebDriverWrapper.sleep;
 
 public class Timer
 {
@@ -72,4 +77,43 @@ public boolean isTimedOut()
     {
         return _cancelled || timeRemaining().isNegative();
     }
+
+    /**
+     * Wait for Supplier to return non-null non-false value. Supplier is guaranteed to run at least once.
+     * @return final result of Supplier.get()
+     */
+    @Contract(pure = true)
+    public <T> T waitFor(Supplier<T> checker)
+    {
+        T result;
+        do
+        {
+            result = checker.get();
+            if (isSuccessResult(result))
+                break;
+            sleep(100);
+        } while (!isTimedOut());
+
+        return result;
+    }
+
+    public <T> T waitFor(Supplier<T> checker, Supplier<String> failMessageSupplier)
+    {
+        T result = waitFor(checker);
+        if (!isSuccessResult(result))
+        {
+            throw new TimeoutException(failMessageSupplier.get() + TestLogger.formatElapsedTime(_timeout.toMillis()));
+        }
+        return result;
+    }
+
+    public <T> T waitFor(Supplier<T> checker, String failMessage)
+    {
+        return waitFor(checker, () -> failMessage);
+    }
+
+    private static <T> boolean isSuccessResult(T result)
+    {
+        return result != null && !Boolean.FALSE.equals(result);
+    }
 }
diff --git a/src/org/labkey/test/util/perf/JsonPerfScenarioHelper.java b/src/org/labkey/test/util/perf/JsonPerfScenarioHelper.java
@@ -9,6 +9,7 @@
 import org.labkey.test.params.perf.PerfScenario;
 import org.labkey.test.stress.AbstractScenario;
 import org.labkey.test.stress.RequestInfoTsvWriter;
+import org.labkey.test.util.LogMethod;
 import org.labkey.test.util.TestDateUtils;
 import org.labkey.test.util.TestLogger;
 import org.labkey.test.util.Timer;
@@ -56,6 +57,7 @@ public JsonPerfScenarioHelper setImportThreads(int importThreads)
         return this;
     }
 
+    @LogMethod
     public Map<String, Result> runPerfScenarios(List<PerfScenario> scenarios) throws Exception
     {
         final ExecutorService importExecutor = Executors.newFixedThreadPool(importThreads);
@@ -94,7 +96,8 @@ private Result startImport(PerfScenario perfScenario)
         }
         ImportDataCommand command = new ImportDataCommand(schemaName, perfScenario.getTypeName());
         command.setFile(perfDataFileSupplier.apply(perfScenario.getFileName()));
-        command.setTimeout(Math.max(connection.getTimeout(), perfScenario.getAverage() * importThreads));
+        command.setTimeout(Math.max(connection.getTimeout(), perfScenario.getAverage() * importThreads) * 2);
+        command.setInsertOption(ImportDataCommand.InsertOption.MERGE);
         Timer timer = new Timer();
         String msgSuffix = "";
         try