[ #154578064 ] - some updates

Author: denshikov-vovan

tests/helpers.ts

@@ -1,10 +1,11 @@
 import * as net from 'net';
 import {Client, ConnectionConfig} from 'pg';
 import {buildLogger, Logger} from '../lib/util';
-import {ConnectionPool, PoolOptions} from '../lib/Pool';
+import {CLOSED_EVENT, ConnectionPool, PoolOptions} from '../lib/Pool';
 import {settings} from './settings';
 
-export const PROXY_SERVER_PORT = 3000;
+export const PROXY_SERVER_PORT   = 3000;
+export const FAKE_PG_SERVER_PORT = 3001;
 
 export function createNewPool(poolOptions: PoolOptions = {}, connectionOptions: ConnectionConfig = {}, logger?: Logger): ConnectionPool {
     const poolSettings = Object.assign({}, settings.pool, poolOptions);
@@ -22,6 +23,39 @@ export function createClient(pool: ConnectionPool): Promise<Client> {
     return new Promise((resolve, reject) => pool.acquire((err, client) => err ? reject(err) : resolve(client)));
 }
 
+export function removeAllClientAndShutdownPool(pool: ConnectionPool): Promise<void> {
+    const removeClient = (client: Client) => {
+        const idleTimeout = (pool as any).idle.get(client);
+
+        if (idleTimeout) {
+            clearTimeout(idleTimeout);
+            (pool as any).idle.delete(client);
+        }
+
+        (pool as any).clients.delete(client);
+
+        if ((client as any).connection) {
+            (client as any).connection.stream.destroy();
+        } else {
+            client.end();
+        }
+    };
+
+    const shutdownPool = (): Promise<void> => {
+        return new Promise((resolve, reject) => {
+            pool.shutdown(error => {
+                error ? reject(error) : resolve();
+            });
+        });
+    };
+
+    for (let client of (pool as any).clients.keys()) {
+        removeClient(client);
+    }
+
+    return shutdownPool();
+}
+
 export function pgProxyServer(timeout: number, cb): net.Server {
     const server = net.createServer(proxySocket => {
         const pgSocket = new net.Socket();
@@ -56,3 +90,11 @@ export function pgProxyServer(timeout: number, cb): net.Server {
 
     return server;
 }
+
+export function fakePgServer(cb: Function): net.Server {
+    const server = net.createServer();
+
+    server.listen(FAKE_PG_SERVER_PORT, cb);
+
+    return server;
+}

tests/poolConnectionTimeout.spec.ts

@@ -1,15 +1,15 @@
 import {expect} from 'chai';
 
 import {ConnectionError} from '../lib/errors';
-import {createNewPool, pgProxyServer, PROXY_SERVER_PORT} from './helpers';
+import {createNewPool, fakePgServer, removeAllClientAndShutdownPool, FAKE_PG_SERVER_PORT} from './helpers';
 
-const connectionOptions = {port: PROXY_SERVER_PORT};
+const connectionOptions = {port: FAKE_PG_SERVER_PORT};
 
 let server;
 
 describe('Pool connection timeout;', () => {
     before(done => {
-        server = pgProxyServer(250, done);
+        server = fakePgServer(done);
     });
 
     after(done => {
@@ -19,17 +19,21 @@ describe('Pool connection timeout;', () => {
     it('should callback with an error if timeout is passed', done => {
         const pool = createNewPool({connectionTimeout: 150}, connectionOptions);
 
-        pool.acquire((err, client) => {
+        pool.acquire(async (err, client) => {
             try {
                 expect(err).to.be.an.instanceof(ConnectionError);
                 expect(err.message).to.contain('Connection request has timed out');
                 expect(pool.idleCount).to.equal(0);
                 expect(pool.totalCount).to.equal(1);
                 expect(client).to.be.undefined;
 
-                pool.shutdown(done);
+                await removeAllClientAndShutdownPool(pool);
+
+                done();
             } catch (err) {
-                pool.shutdown(() => done(err));
+                removeAllClientAndShutdownPool(pool)
+                    .then( () => done(err))
+                    .catch(() => done(err));
             }
         });
     });
@@ -58,27 +62,33 @@ describe('Pool connection timeout;', () => {
             expect(pool.idleCount).to.equal(0);
             expect(pool.totalCount).to.equal(15);
 
-            pool.shutdown(done);
+            await removeAllClientAndShutdownPool(pool);
+
+            done();
         } catch (err) {
-            pool.shutdown(() => done(err));
+            removeAllClientAndShutdownPool(pool)
+                .then( () => done(err))
+                .catch(() => done(err));
         }
     });
 
     it('should timeout on checkout of used connection', done => {
-        const pool = createNewPool({connectionTimeout: 400, maxSize: 1}, connectionOptions);
+        const pool = createNewPool({connectionTimeout: 400, maxSize: 1});
+
         try {
             pool.acquire((err, client) => {
                 expect(err).to.be.undefined;
                 expect(client).to.not.be.undefined;
                 expect(pool.totalCount).to.equal(1);
 
-                pool.acquire((err, client) => {
+                pool.acquire(async (err, client) => {
                     expect(err).to.be.an.instanceof(ConnectionError);
                     expect(err.message).to.contain('Connection request has timed out');
                     expect(client).to.be.undefined;
                     expect(pool.totalCount).to.equal(1);
 
                     (pool as any).clients.entries().next().value[0].release();
+
                     pool.shutdown(done);
                 });
             });
@@ -88,7 +98,7 @@ describe('Pool connection timeout;', () => {
     });
 
     it('should timeout on query if all clients are busy', done => {
-        const pool = createNewPool({connectionTimeout: 400, maxSize: 1}, connectionOptions);
+        const pool = createNewPool({connectionTimeout: 400, maxSize: 1});
 
         pool.acquire((err, client) => {
             expect(err).to.be.undefined;
@@ -109,7 +119,7 @@ describe('Pool connection timeout;', () => {
     });
 
     it('should recover from timeout errors', done => {
-        const pool = createNewPool({connectionTimeout: 400, maxSize: 1}, connectionOptions);
+        const pool = createNewPool({connectionTimeout: 400, maxSize: 1});
 
         pool.acquire((err, client) => {
             expect(err).to.be.undefined;