From 138071f44b595e0f0453af359dd6f5a06c6469e9 Mon Sep 17 00:00:00 2001
From: labkey-tchad <tchad@labkey.com>
Date: Mon, 15 Jul 2024 12:41:53 -0700
Subject: [PATCH] Allow custom connection in AuditLogHelper

---
 src/org/labkey/test/util/AuditLogHelper.java | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/org/labkey/test/util/AuditLogHelper.java b/src/org/labkey/test/util/AuditLogHelper.java
index 7046b72a09..a4a12ed536 100644
--- a/src/org/labkey/test/util/AuditLogHelper.java
+++ b/src/org/labkey/test/util/AuditLogHelper.java
@@ -1,6 +1,7 @@
 package org.labkey.test.util;
 
 import org.labkey.remoteapi.CommandException;
+import org.labkey.remoteapi.Connection;
 import org.labkey.remoteapi.query.ContainerFilter;
 import org.labkey.remoteapi.query.SelectRowsCommand;
 import org.labkey.remoteapi.query.SelectRowsResponse;
@@ -12,14 +13,22 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 public class AuditLogHelper
 {
     private final LabKeySiteWrapper _wrapper;
+    private final ConnectionSupplier _connectionSupplier;
 
-    public AuditLogHelper(LabKeySiteWrapper wrapper)
+    public AuditLogHelper(LabKeySiteWrapper wrapper, ConnectionSupplier connectionSupplier)
     {
         _wrapper = wrapper;
+        _connectionSupplier = connectionSupplier;
+    }
+
+    public AuditLogHelper(LabKeySiteWrapper wrapper)
+    {
+        this(wrapper, wrapper::createDefaultConnection);
     }
 
     public Integer getLatestAuditRowId(String auditTable) throws IOException, CommandException
@@ -32,11 +41,11 @@ public Integer getLatestAuditRowId(String auditTable) throws IOException, Comman
         selectRows.setMaxRows(1);
         selectRows.setContainerFilter(ContainerFilter.AllFolders);
 
-        SelectRowsResponse response = selectRows.execute(_wrapper.createDefaultConnection(), null);
+        SelectRowsResponse response = selectRows.execute(_connectionSupplier.get(), null);
         List<Map<String, Object>> rows = response.getRows();
         if (rows.isEmpty())
         {
-            return -1;
+            return 0;
         }
         return (Integer) rows.get(0).get(rowId);
     }
@@ -59,4 +68,9 @@ public DataRegionTable goToAuditEventView(String eventType)
         }
         return new DataRegionTable("query", _wrapper);
     }
+
+    public interface ConnectionSupplier
+    {
+        Connection get() throws IOException, CommandException;
+    }
 }