From b7cb2c98d7d24ae6bc577aff54c602612dcbbc48 Mon Sep 17 00:00:00 2001
From: Ugaitz Urien <ugaitz.urien@datadoghq.com>
Date: Thu, 1 Aug 2024 15:58:09 +0200
Subject: [PATCH] Adding instrumentation tests

---
 packages/datadog-instrumentations/src/pg.js   |   2 +-
 .../datadog-instrumentations/test/pg.spec.js  | 112 ++++++++++++++++++
 2 files changed, 113 insertions(+), 1 deletion(-)
 create mode 100644 packages/datadog-instrumentations/test/pg.spec.js

diff --git a/packages/datadog-instrumentations/src/pg.js b/packages/datadog-instrumentations/src/pg.js
index 547ff3aa66d..5358c6e0c5b 100644
--- a/packages/datadog-instrumentations/src/pg.js
+++ b/packages/datadog-instrumentations/src/pg.js
@@ -72,7 +72,7 @@ function wrapQuery (query) {
       if (abortController.signal.aborted) {
         const error = abortController.signal.reason || new Error('Aborted')
         const reusingQuery = typeof pgQuery.submit === 'function'
-        const callback = (reusingQuery && pgQuery.callback) || typeof arguments[arguments.length - 1]
+        const callback = (reusingQuery && pgQuery.callback) || arguments[arguments.length - 1]
 
         finish(error)
 
diff --git a/packages/datadog-instrumentations/test/pg.spec.js b/packages/datadog-instrumentations/test/pg.spec.js
new file mode 100644
index 00000000000..e6d5c56a51c
--- /dev/null
+++ b/packages/datadog-instrumentations/test/pg.spec.js
@@ -0,0 +1,112 @@
+const agent = require('../../dd-trace/test/plugins/agent')
+const dc = require('dc-polyfill')
+
+const clients = {
+  pg: pg => pg.Client
+}
+
+if (process.env.PG_TEST_NATIVE === 'true') {
+  clients['pg.native'] = pg => pg.native.Client
+}
+
+describe('pg instrumentation', () => {
+  withVersions('pg', 'pg', version => {
+    const queryStartChannel = dc.channel('apm:pg:query:start')
+
+    let pg
+    let client
+
+    function abortQuery({ abortController }) {
+      const error = new Error('Test')
+      abortController.abort(error)
+
+      if (!abortController.signal.reason) {
+        abortController.signal.reason = error
+      }
+    }
+
+    before(() => {
+      return agent.load(['pg'])
+    })
+
+    Object.keys(clients).forEach(implementation => {
+      beforeEach(done => {
+        pg = require(`../../../versions/pg@${version}`).get()
+        const Client = clients[implementation](pg)
+
+        client = new Client({
+          host: '127.0.0.1',
+          user: 'postgres',
+          password: 'postgres',
+          database: 'postgres',
+          application_name: 'test'
+        })
+
+        client.connect(err => done(err))
+      })
+
+      // afterEach((done) => {
+      //   client.end((err) => {
+      //     done(err)
+      //   })
+      // })
+
+      describe('pg.Client', () => {
+        describe('abortController', () => {
+          afterEach(() => {
+            if (queryStartChannel.hasSubscribers) {
+              queryStartChannel.unsubscribe(abortQuery)
+            }
+          })
+
+          describe('using callback', () => {
+            it('Should not fail if it is not aborted', (done) => {
+              client.query('SELECT 1', (err) => {
+                client.end()
+                done(err)
+              })
+            })
+
+            it('Should abort query', (done) => {
+              queryStartChannel.subscribe(abortQuery)
+
+              client.query('SELECT 1', (err) => {
+                try {
+                  client.end()
+                } catch (e) {
+                  console.error(e)
+                }
+                if (err && err.message === 'Test') {
+                  return done()
+                }
+
+                done(new Error('Query was not aborted'))
+              })
+            })
+          })
+
+          describe('using promise', () => {
+            it('Should not fail if it is not aborted', async () => {
+              await client.query('SELECT 1')
+              client.end()
+            })
+
+            it('Should abort query', async () => {
+              queryStartChannel.subscribe(abortQuery)
+
+              try {
+                await client.query('SELECT 1')
+
+                throw new Error('Query was not aborted')
+              } catch (err) {
+                if (!err || err.message !== 'Test') {
+                }
+                  throw err
+              }
+            })
+          })
+        })
+      })
+    })
+  })
+})