OpenTelemetry Instrumentation - Can't get my tests to work - NodeJS

231 Views Asked by At

I'm currently in the course of writing an easy-to-use instrumentation wrapper library using Openetelemetry and NodeJS.

For that, I'm heavily using the node-auto-instrumentation package but I do have some more specific packages in use, which I also want to auto-instrument. In particular I'm talking about the package "pg-copy-streams".

The actual Instrumentation Class for that package, I came up with, works absolutely fine. However, since I wanted to open-source the package, I also want to write tests for it ... but really cannot get any of them to work. It even fails at the bare minimum "isWrapped"

Since the instrumentation code is no "secret" at all, I can show you, what I came up with:

import { InstrumentationBase, InstrumentationNodeModuleDefinition, isWrapped } from '@opentelemetry/instrumentation';
import { diag, SpanKind, SpanStatusCode } from '@opentelemetry/api';
import * as copy from 'pg-copy-streams';

export class PgCopyStreamsInstrumentation extends InstrumentationBase<typeof copy> {
    constructor() {
        super('opentelemetry-instrumentation-pg-copy-streams', '1.0.0');
    }

    init(): InstrumentationNodeModuleDefinition<typeof copy>[] {
        return [
            new InstrumentationNodeModuleDefinition<typeof copy>(
                'pg-copy-streams',
                ['*'],
                moduleExports => {
                    this._wrap(moduleExports, 'from', this._patch('pg-copy-streams.from').bind(this));
                    this._wrap(moduleExports, 'to', this._patch('pg-copy-streams.to').bind(this));
                    diag.info(`opentelemetry-instrumentation-pg-copy-streams initialized`);
                    return moduleExports;
                },
                moduleExports => {
                    if (isWrapped(moduleExports.from)) {
                        this._unwrap(moduleExports, 'from');
                    }
                    if (isWrapped(moduleExports.to)) {
                        this._unwrap(moduleExports, 'to');
                    }
                }
            )
        ];
    }

    private _patch(operation: string) {
        return (original: (...args: unknown[]) => unknown) => {
            return (sql: string, ...args: unknown[]) => {
                const span = this.tracer.startSpan(operation, {
                    kind: SpanKind.CLIENT,
                    attributes: {
                        'copy.db.statement': sql
                    }
                });

                const returnedStream = original.apply(this, [sql, ...args]);

                let rowsProcessed = 0;
                let bytesTransferred = 0;
                const startTime = Date.now();

                returnedStream.on('data', (chunk: Buffer) => {
                    rowsProcessed += (chunk.toString().match(/\n/g) || []).length;
                    bytesTransferred += chunk.length;
                });

                returnedStream.on('end', () => {
                    const endTime = Date.now();
                    const durationSeconds = (endTime - startTime) / 1000;
                    const dataTransferRateBytesPerSec = bytesTransferred / durationSeconds; 
                    span.setAttribute('copy.rows_processed', rowsProcessed);
                    span.setAttribute(
                        'copy.data_transfer_rate_bps',
                        parseFloat(dataTransferRateBytesPerSec.toFixed(2))
                    );
                    span.setStatus({ code: SpanStatusCode.OK });
                    span.end();
                });

                returnedStream.on('error', (error: Error) => {
                    span.setStatus({ code: SpanStatusCode.ERROR, message: error.message });
                    span.end();
                });

                return returnedStream;
            };
        };
    }
}

const pgCopyStreamsInstrumentation = new PgCopyStreamsInstrumentation();
export default pgCopyStreamsInstrumentation;

And here is my started, not working, test cases file

import { PgCopyStreamsInstrumentation } from '../src/custom-instrumentations/pg-copy-streams/pg-copy-streams';

import { isWrapped } from "@opentelemetry/instrumentation";
import { getTestSpans, registerInstrumentationTesting } from '@opentelemetry/contrib-test-utils';

const instrumentation = registerInstrumentationTesting(new PgCopyStreamsInstrumentation());
import * as copy from 'pg-copy-streams';

describe('PgCopyStreamsInstrumentation', () => {

    beforeEach(() => {
        instrumentation.enable();
    });

    afterEach(() => {
        instrumentation.disable();
    });

    it('should create a span for the from function', () => {
        expect(isWrapped(copy.from)).toBe(true);

        copy.from('SELECT * FROM users');

        const spans = getTestSpans();
        expect(spans).toHaveLength(1);
        expect(spans[0].name).toEqual('pg-copy-streams.from');
    });

});

As said ... the actual code works perfectly fine and integrates with my tracer setup, creates spans and traces ... it's just that I cannot get it to run in the isolated (Jest) test environment ... any hint would be highly appreciated since I struggle with that for quite some hours already :)

0

There are 0 best solutions below