feat: times (#80)

traceloop · Jun 20, 2023 · 4fc74e4 · 4fc74e4
1 parent 81e4ee5
commit 4fc74e4
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 57 deletions.
diff --git a/docs/jest-otel/syntax/db-redis.mdx b/docs/jest-otel/syntax/db-redis.mdx
@@ -28,7 +28,7 @@ So, a complete assertion can look like:
 
 ```js
 expectTrace(traceloop.serviceByName('redis-service'))
-  .toSendRedisCommend()
+  .toSendRedisCommend({ times: 2 }) // optional times parameter (defaults to one)
   .withDatabaseName('redis-db')
   .withStatement(
     /^HGET/

diff --git a/packages/expect-opentelemetry/src/index.ts b/packages/expect-opentelemetry/src/index.ts
@@ -29,13 +29,17 @@ const serviceMatchers = {
   toSendRedisCommand,
 };
 
+interface MatcherOptions {
+  times: number;
+}
+
 interface TraceMatchers {
   toReceiveHttpRequest(): HttpRequest;
   toSendHttpRequest(): HttpRequest;
-  toQueryPostgreSQL(): PostgreSQLQuery;
+  toQueryPostgreSQL(options?: MatcherOptions): PostgreSQLQuery;
   toReceiveGrpcRequest(): GrpcRequest;
   toSendGrpcRequest(): GrpcRequest;
-  toSendRedisCommand(): RedisCommand;
+  toSendRedisCommand(options?: MatcherOptions): RedisCommand;
 }
 
 function createMatcher(matcher, type) {

diff --git a/packages/expect-opentelemetry/src/matchers/service/to-query-postgresql.ts b/packages/expect-opentelemetry/src/matchers/service/to-query-postgresql.ts
@@ -3,25 +3,32 @@ import { Service } from '../../resources/service';
 import { opentelemetry } from '@traceloop/otel-proto';
 import { PostgreSQLQuery } from '../../resources/postgresql-query';
 
-export function toQueryPostgreSQL(service: Service): PostgreSQLQuery {
+export function toQueryPostgreSQL(
+  service: Service,
+  options = { times: 1 },
+): PostgreSQLQuery {
   const { name: serviceName, spans } = service;
 
   const filteredSpans = spans.filter((span) => {
     return (
       span.kind ===
         opentelemetry.proto.trace.v1.Span.SpanKind.SPAN_KIND_CLIENT &&
-      span.attributes?.find((attribute) => {
-        return (
-          attribute.key === SemanticAttributes.DB_SYSTEM &&
-          attribute.value?.stringValue === 'postgresql'
-        );
-      })
+      span.attributes?.find(
+        (attribute: opentelemetry.proto.common.v1.IKeyValue) => {
+          return (
+            attribute.key === SemanticAttributes.DB_SYSTEM &&
+            attribute.value?.stringValue === 'postgresql'
+          );
+        },
+      )
     );
   });
 
-  if (filteredSpans.length === 0) {
-    throw new Error(`No query by ${serviceName} to postgresql was found`);
+  if (filteredSpans.length < options.times) {
+    throw new Error(
+      `Expected ${options.times} queries by ${serviceName} to postgresql, but found ${filteredSpans.length}.`,
+    );
   }
 
-  return new PostgreSQLQuery(filteredSpans, serviceName);
+  return new PostgreSQLQuery(filteredSpans, serviceName, options.times);
 }
diff --git a/packages/expect-opentelemetry/src/matchers/service/to-send-redis-command.ts b/packages/expect-opentelemetry/src/matchers/service/to-send-redis-command.ts
@@ -2,7 +2,10 @@ import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
 import { opentelemetry } from '@traceloop/otel-proto';
 import { RedisCommand, Service } from '../../resources';
 
-export function toSendRedisCommand(service: Service): RedisCommand {
+export function toSendRedisCommand(
+  service: Service,
+  options = { times: 1 },
+): RedisCommand {
   const { name: serviceName, spans } = service;
 
   const filteredSpans = spans.filter((span) => {
@@ -20,9 +23,11 @@ export function toSendRedisCommand(service: Service): RedisCommand {
     );
   });
 
-  if (filteredSpans.length === 0) {
-    throw new Error(`No redis command from ${serviceName} found`);
+  if (filteredSpans.length < options.times) {
+    throw new Error(
+      `Expected ${options.times} queries by ${serviceName} to redis, but found ${filteredSpans.length}.`,
+    );
   }
 
-  return new RedisCommand(filteredSpans, serviceName);
+  return new RedisCommand(filteredSpans, serviceName, options.times);
 }
diff --git a/packages/expect-opentelemetry/src/resources/postgresql-query.ts b/packages/expect-opentelemetry/src/resources/postgresql-query.ts
@@ -12,6 +12,7 @@ export class PostgreSQLQuery {
   constructor(
     readonly spans: opentelemetry.proto.trace.v1.ISpan[],
     private readonly serviceName: string,
+    private readonly times = 1,
   ) {}
 
   withDatabaseName(name: string | RegExp, options?: CompareOptions) {
@@ -22,18 +23,19 @@ export class PostgreSQLQuery {
       options,
     );
 
-    if (filteredSpans.length === 0) {
-      throw new Error(
-        `No query by ${
-          this.serviceName
-        } to postgresql with database name ${name} was found. Found db names: ${extractAttributeStringValues(
-          this.spans,
-          SemanticAttributes.DB_NAME,
-        )}`,
-      );
+    if (filteredSpans.length < this.times) {
+      throw new Error(`Expected ${this.times} queries by ${
+        this.serviceName
+      } to postgresql with database name ${name}, but found ${
+        filteredSpans.length
+      }.\n
+      Found db names:\n ${extractAttributeStringValues(
+        this.spans,
+        SemanticAttributes.DB_NAME,
+      )}`);
     }
 
-    return new PostgreSQLQuery(filteredSpans, this.serviceName);
+    return new PostgreSQLQuery(filteredSpans, this.serviceName, this.times);
   }
 
   withStatement(statement: string | RegExp, options?: CompareOptions) {
@@ -44,18 +46,16 @@ export class PostgreSQLQuery {
       options,
     );
 
-    if (filteredSpans.length === 0) {
-      throw new Error(
-        `No query by ${
-          this.serviceName
-        } to postgresql with statement ${statement} was found. Found statements:\n ${extractAttributeStringValues(
-          this.spans,
-          SemanticAttributes.DB_STATEMENT,
-        ).join('\n\n')}`,
-      );
+    if (filteredSpans.length < this.times) {
+      throw new Error(`Expected ${this.times} queries by ${
+        this.serviceName
+      } to postgresql with statement ${statement}, but found ${
+        filteredSpans.length
+      }.\n
+      Found statements:\n${printStatements(this.spans)}`);
     }
 
-    return new PostgreSQLQuery(filteredSpans, this.serviceName);
+    return new PostgreSQLQuery(filteredSpans, this.serviceName, this.times);
   }
 
   withOperations(...operations: string[]) {
@@ -76,18 +76,14 @@ export class PostgreSQLQuery {
       );
     });
 
-    if (filteredSpans.length === 0) {
+    if (filteredSpans.length < this.times) {
       throw new Error(
-        `No query by ${
-          this.serviceName
-        } to postgresql with operations ${operations} was found. Found statements: ${extractAttributeStringValues(
-          this.spans,
-          SemanticAttributes.DB_STATEMENT,
-        ).join('\n\n')}`,
+        `Expected ${this.times} queries by ${this.serviceName} to postgresql with operations ${operations}, but found ${filteredSpans.length}.\n` +
+          `Found statements:\n${printStatements(this.spans)}`,
       );
     }
 
-    return new PostgreSQLQuery(filteredSpans, this.serviceName);
+    return new PostgreSQLQuery(filteredSpans, this.serviceName, this.times);
   }
 
   withTables(...tables: string[]) {
@@ -117,17 +113,25 @@ export class PostgreSQLQuery {
       );
     });
 
-    if (filteredSpans.length === 0) {
+    if (filteredSpans.length < this.times) {
       throw new Error(
-        `No query by ${
-          this.serviceName
-        } to postgresql with tables ${tables} was found. Found statements: ${extractAttributeStringValues(
-          this.spans,
-          SemanticAttributes.DB_STATEMENT,
-        ).join('\n\n')}`,
+        `Expected ${this.times} queries by ${this.serviceName} to postgresql with tables ${tables}, but found ${filteredSpans.length}.\n` +
+          `Found statements:\n${printStatements(this.spans)}`,
       );
     }
 
-    return new PostgreSQLQuery(filteredSpans, this.serviceName);
+    return new PostgreSQLQuery(filteredSpans, this.serviceName, this.times);
   }
 }
+
+const printStatements = (spans: opentelemetry.proto.trace.v1.ISpan[]) => {
+  const MAX_LEN = 100;
+  return extractAttributeStringValues(spans, SemanticAttributes.DB_STATEMENT)
+    .map((statement) => {
+      if (statement.length > MAX_LEN) {
+        return `${statement.slice(0, MAX_LEN)}...`;
+      }
+      return statement;
+    })
+    .join('\n');
+};
diff --git a/packages/expect-opentelemetry/src/resources/redis-command.ts b/packages/expect-opentelemetry/src/resources/redis-command.ts
@@ -9,6 +9,7 @@ export class RedisCommand {
   constructor(
     readonly spans: opentelemetry.proto.trace.v1.ISpan[],
     private readonly serviceName: string,
+    private readonly times = 1,
   ) {}
 
   withDatabaseName(name: string | RegExp, options: CompareOptions) {
@@ -19,9 +20,9 @@ export class RedisCommand {
       options,
     );
 
-    if (filteredSpans.length === 0) {
+    if (filteredSpans.length < this.times) {
       throw new Error(
-        `No redis command from service ${this.serviceName} to database ${name} found`,
+        `Expected ${this.times} queries by ${this.serviceName} to redis with database name ${name}, but found ${filteredSpans.length}.`,
       );
     }
 
@@ -36,9 +37,9 @@ export class RedisCommand {
       options,
     );
 
-    if (filteredSpans.length === 0) {
+    if (filteredSpans.length < this.times) {
       throw new Error(
-        `No redis command with statement ${statement} from service ${this.serviceName} found`,
+        `Expected ${this.times} queries by ${this.serviceName} to redis with statement ${statement}, but found ${filteredSpans.length}.`,
       );
     }