From c90d5ed43b7380405941da01193f8c1c0154cb20 Mon Sep 17 00:00:00 2001
From: Weimin Yu <weiminyu@google.com>
Date: Tue, 14 Jan 2020 10:59:55 -0500
Subject: [PATCH] Fix flaky SQL integration tests (#441)

* Fix flaky SQL integration tests

As a sanity check, the text fixture verifies that all test jdbc connections
have been closed. However, the query to get the connection count is too
general, and may count background admin tasks such as autovacuum and
replication etc, that may appear at any time.
---
 .../JpaTransactionManagerRule.java            | 49 ++++++++++---------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
index 4495795a6..bfc96d977 100644
--- a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
+++ b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
@@ -66,6 +66,9 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
       "google/registry/model/transaction/cleanup_database.sql";
   private static final String MANAGEMENT_DB_NAME = "management";
   private static final String POSTGRES_DB_NAME = "postgres";
+  // The type of JDBC connections started by the tests. This string value
+  // is documented in PSQL's official user guide.
+  private static final String CONNECTION_BACKEND_TYPE = "client backend";
 
   private final DateTime now = DateTime.now(UTC);
   private final FakeClock clock = new FakeClock(now);
@@ -74,9 +77,6 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
   private final ImmutableMap userProperties;
 
   private static final JdbcDatabaseContainer database = create();
-  private static final long ACTIVE_CONNECTIONS_BASELINE =
-      getActiveConnectionCountByUser(database.getUsername());
-  ;
   private static final HibernateSchemaExporter exporter =
       HibernateSchemaExporter.create(
           database.getJdbcUrl(), database.getUsername(), database.getPassword());
@@ -123,7 +123,7 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
       builder.put(Environment.SHOW_SQL, "true");
       properties = builder.build();
     }
-    assertNormalActiveConnection();
+    assertNoActiveTestConnections();
     emf =
         createEntityManagerFactory(
             getJdbcUrlFor(POSTGRES_DB_NAME),
@@ -144,29 +144,32 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
       emf = null;
     }
     cachedTm = null;
-    assertNormalActiveConnection();
-  }
-
-  private static long getActiveConnectionCountByUser(String userName) {
-    try (Connection conn = createConnection(POSTGRES_DB_NAME);
-        Statement statement = conn.createStatement()) {
-      ResultSet rs =
-          statement.executeQuery(
-              "SELECT COUNT(1) FROM pg_stat_activity WHERE usename = '" + userName + "'");
-      rs.next();
-      return rs.getLong(1);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+    assertNoActiveTestConnections();
   }
 
   /**
-   * This function throws exception if it detects connection leak by checking the metadata table
-   * pg_stat_activity.
+   * Asserts that all JDBC connections started by the test classes to the 'postgres' database have
+   * been closed.
    */
-  private void assertNormalActiveConnection() {
-    assertThat(getActiveConnectionCountByUser(database.getUsername()))
-        .isEqualTo(ACTIVE_CONNECTIONS_BASELINE);
+  private void assertNoActiveTestConnections() {
+    // Use the 'management' db to connect so that this connection needs not to be accounted for.
+    try (Connection conn = createConnection(MANAGEMENT_DB_NAME);
+        Statement statement = conn.createStatement()) {
+      // Note: Since we use the admin user (returned by container's getUserName() method)
+      // in tests, we need to filter connections by database name and/or backend type to filter out
+      // connections for management tasks.
+      ResultSet rs =
+          statement.executeQuery(
+              String.format(
+                  "SELECT COUNT(1) FROM pg_stat_activity WHERE usename = '%1s'"
+                      + " and datname = '%2s' "
+                      + " and backend_type = '%3s'",
+                  database.getUsername(), POSTGRES_DB_NAME, CONNECTION_BACKEND_TYPE));
+      rs.next();
+      assertThat(rs.getLong(1)).isEqualTo(0);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
   }
 
   private static String readSqlInClassPath(String sqlScriptPath) {