From 59fca7845d0891bb79a994465a4cd0e1795b875d Mon Sep 17 00:00:00 2001
From: Nenad Noveljic <18366081+nenadnoveljic@users.noreply.github.com>
Date: Sat, 13 Jul 2024 21:39:20 +0200
Subject: [PATCH] Test classes

---
 .../jdbc/src/test/groovy/test/TestConnection.groovy   | 10 +++++++++-
 .../src/test/groovy/test/TestDatabaseMetaData.groovy  | 11 ++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy
index defb43d9792..9e75ef23d26 100644
--- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy
+++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy
@@ -28,6 +28,7 @@ class TestConnection implements Connection {
     }
   }
 
+  private DatabaseMetaData metadata
 
   @Override
   Statement createStatement() throws SQLException {
@@ -77,7 +78,14 @@ class TestConnection implements Connection {
 
   @Override
   DatabaseMetaData getMetaData() throws SQLException {
-    return new TestDatabaseMetaData()
+    if (metadata == null) {
+      return new TestDatabaseMetaData()
+    }
+    return metadata
+  }
+
+  void setMetaData(DatabaseMetaData metadata) {
+    this.metadata = metadata
   }
 
   @Override
diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestDatabaseMetaData.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestDatabaseMetaData.groovy
index 198f1832dd2..ce5b3e783c7 100644
--- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestDatabaseMetaData.groovy
+++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestDatabaseMetaData.groovy
@@ -12,6 +12,8 @@ class TestDatabaseMetaData implements DatabaseMetaData {
     return false
   }
 
+  private String url
+
   @Override
   boolean allTablesAreSelectable() throws SQLException {
     return false
@@ -19,7 +21,14 @@ class TestDatabaseMetaData implements DatabaseMetaData {
 
   @Override
   String getURL() throws SQLException {
-    return "jdbc:testdb://localhost"
+    if (url == null) {
+      return "jdbc:testdb://localhost"
+    }
+    return url
+  }
+
+  void setURL(String url) {
+    this.url = url
   }
 
   @Override