add more tests, and some fixes to go along with them

src/config.rs

@@ -23,7 +23,7 @@ macro_rules! get_arr {
     };
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum InstrumentationOperator {
     Callback,
     Promise,
@@ -52,6 +52,7 @@ impl InstrumentationOperator {
     }
 }
 
+#[derive(Debug)]
 pub struct InstrumentationConfig {
     pub module_name: String,
     pub version_range: Range,

src/function_query.rs

@@ -10,6 +10,7 @@ macro_rules! get_str {
     };
 }
 
+#[derive(Debug)]
 pub enum FunctionType {
     FunctionDeclaration,
     FunctionExpression,
@@ -27,6 +28,7 @@ impl FunctionType {
     }
 }
 
+#[derive(Debug)]
 pub enum FunctionKind {
     Sync,
     Async,
@@ -63,6 +65,7 @@ impl FunctionKind {
     }
 }
 
+#[derive(Debug)]
 pub struct FunctionQuery {
     pub name: String,
     pub class: Option<String>,

src/instrumentation.rs

@@ -38,6 +38,11 @@ impl Instrumentation {
         }
     }
 
+    pub(crate) fn reset(&mut self) {
+        self.count = 0;
+        self.is_correct_class = false;
+    }
+
     fn new_fn(&self, body: BlockStmt) -> ArrowExpr {
         ArrowExpr {
             params: vec![],
@@ -70,7 +75,9 @@ impl Instrumentation {
         define_channel
     }
 
-    fn insert_tracing(&self, body: &mut BlockStmt) {
+    fn insert_tracing(&mut self, body: &mut BlockStmt) {
+        self.count += 1;
+
         let original_stmts = std::mem::take(&mut body.stmts);
 
         // Create a new BlockStmt with the original statements
@@ -102,7 +109,9 @@ impl Instrumentation {
         ];
     }
 
-    fn insert_constructor_tracing(&self, body: &mut BlockStmt) {
+    fn insert_constructor_tracing(&mut self, body: &mut BlockStmt) {
+        self.count += 1;
+
         let original_stmts = std::mem::take(&mut body.stmts);
 
         let ch_ident = ident!(format!("tr_ch_apm${}", &self.config.channel_name));

src/lib.rs

@@ -104,7 +104,8 @@ macro_rules! visit_with_all {
     ($self:expr, $method:ident, $item:expr) => {
         let mut recurse = false;
         for instr in &mut $self.instrumentations {
-            recurse = recurse || instr.$method($item);
+            let needs_recurse = instr.$method($item);
+            recurse = recurse || needs_recurse;
         }
         if recurse {
             $item.visit_mut_children_with($self);
@@ -132,6 +133,9 @@ impl VisitMut for InstrumentationVisitor<'_> {
             }
         }
         visit_with_all!(self, visit_mut_module, item);
+        for instr in &mut self.instrumentations {
+            instr.reset();
+        }
     }
 
     fn visit_mut_script(&mut self, item: &mut Script) {
@@ -141,6 +145,9 @@ impl VisitMut for InstrumentationVisitor<'_> {
         );
         item.body.insert(get_script_start_index(item), import);
         visit_with_all!(self, visit_mut_script, item);
+        for instr in &mut self.instrumentations {
+            instr.reset();
+        }
     }
 
     visit_with_all_fn!(visit_mut_fn_decl, FnDecl);

tests/index_cjs/instrumentations.yml

@@ -0,0 +1,13 @@
+version: 1
+instrumentations:
+  - module_name: undici
+    version_range: ">=0.0.1"
+    file_path: index.mjs
+    function_query:
+      class: Undici
+      name: fetch
+      type: method
+      kind: async
+      index: 2
+    operator: tracePromise
+    channel_name: Undici_fetch

tests/index_cjs/mod.js

@@ -0,0 +1,45 @@
+{ 
+  class Undici {
+    async fetch (url) {
+      return 0;
+    }
+  }
+
+  exports.Undici0 = Undici;
+}
+{ 
+  class Undici {
+    async fetch (url) {
+      return 1;
+    }
+  }
+
+  exports.Undici1 = Undici;
+}
+{ 
+  class Undici {
+    async fetch (url) {
+      return 2;
+    }
+  }
+
+  exports.Undici2 = Undici;
+}
+{ 
+  class Undici {
+    async fetch (url) {
+      return 3;
+    }
+  }
+
+  exports.Undici3 = Undici;
+}
+{ 
+  class Undici {
+    async fetch (url) {
+      return 4;
+    }
+  }
+
+  exports.Undici4 = Undici;
+}

tests/index_cjs/test.js

@@ -0,0 +1,27 @@
+const undicis = require('./instrumented.js');
+const { assert, getContext } = require('../common/preamble.js');
+const context = getContext('orchestrion:undici:Undici_fetch');
+
+async function testOne(Undici, num, expectedCtx) {
+  const undici = new Undici;
+  const result = await undici.fetch('https://example.com');
+  assert.strictEqual(result, num);
+  assert.deepStrictEqual(context, expectedCtx);
+  delete context.start;
+  delete context.end;
+  delete context.asyncStart;
+  delete context.asyncEnd;
+}
+
+(async () => {
+  await testOne(undicis.Undici0, 0, {});
+  await testOne(undicis.Undici1, 1, {});
+  await testOne(undicis.Undici2, 2, {
+    start: true,
+    end: true,
+    asyncStart: 2,
+    asyncEnd: 2
+  });
+  await testOne(undicis.Undici3, 3, {});
+  await testOne(undicis.Undici4, 4, {});
+})();

tests/instrumentor_test.rs

@@ -13,6 +13,9 @@ macro_rules! make_test {
                 instrumentor.get_matching_instrumentations("undici", "0.0.1", &file_path);
 
             transpile_and_test(stringify!($name), $mjs, &mut instrumentations);
+
+            // It has to work twice, since we might use the same instrumentor on multiple files
+            transpile_and_test(stringify!($name), $mjs, &mut instrumentations);
         }
     };
 }
@@ -27,6 +30,8 @@ make_test!(expr_cjs, false);
 
 make_test!(class_method_cjs, false);
 
+make_test!(multiple_class_method_cjs, false);
+
 make_test!(object_method_cjs, false);
 
 make_test!(constructor_cjs, false);
@@ -36,3 +41,7 @@ make_test!(constructor_mjs, true);
 make_test!(polyfill_mjs, true);
 
 make_test!(polyfill_cjs, false);
+
+make_test!(index_cjs, false);
+
+make_test!(no_index_cjs, false);

tests/multiple_class_method_cjs/instrumentations.yml

@@ -0,0 +1,22 @@
+version: 1
+instrumentations:
+  - module_name: undici
+    version_range: ">=0.0.1"
+    file_path: index.mjs
+    function_query:
+      class: Undici
+      name: fetch1
+      type: method
+      kind: async
+    operator: tracePromise
+    channel_name: Undici_fetch1
+  - module_name: undici
+    version_range: ">=0.0.1"
+    file_path: index.mjs
+    function_query:
+      class: Undici
+      name: fetch2
+      type: method
+      kind: async
+    operator: tracePromise
+    channel_name: Undici_fetch2

tests/multiple_class_method_cjs/mod.js

@@ -0,0 +1,11 @@
+class Undici {
+    async fetch1 (url) {
+        return 42;
+    }
+
+    async fetch2 (url) {
+        return 43;
+    }
+}
+
+module.exports = Undici;

tests/multiple_class_method_cjs/test.js

@@ -0,0 +1,24 @@
+const Undici = require('./instrumented.js');
+const { assert, getContext } = require('../common/preamble.js');
+const context1 = getContext('orchestrion:undici:Undici_fetch1');
+const context2 = getContext('orchestrion:undici:Undici_fetch2');
+
+(async () => {
+  const undici = new Undici;
+  const result1 = await undici.fetch1('https://example.com');
+  assert.strictEqual(result1, 42);
+  assert.deepStrictEqual(context1, {
+    start: true,
+    end: true,
+    asyncStart: 42,
+    asyncEnd: 42
+  });
+  const result2 = await undici.fetch2('https://example.com');
+  assert.strictEqual(result2, 43);
+  assert.deepStrictEqual(context2, {
+    start: true,
+    end: true,
+    asyncStart: 43,
+    asyncEnd: 43
+  });
+})();

tests/no_index_cjs/instrumentations.yml

@@ -0,0 +1,12 @@
+version: 1
+instrumentations:
+  - module_name: undici
+    version_range: ">=0.0.1"
+    file_path: index.mjs
+    function_query:
+      name: fetch
+      type: decl
+      kind: async
+      # intentionally omitting index
+    operator: tracePromise
+    channel_name: fetch_no_index

tests/no_index_cjs/mod.js

@@ -0,0 +1,5 @@
+async function fetch(url) {
+  return 42;
+}
+
+module.exports = { fetch };

tests/no_index_cjs/test.js

@@ -0,0 +1,13 @@
+const { fetch } = require('./instrumented.js');
+const { assert, getContext } = require('../common/preamble.js');
+const context = getContext('orchestrion:undici:fetch_no_index');
+(async () => {
+  const result = await fetch('https://example.com');
+  assert.strictEqual(result, 42);
+  assert.deepStrictEqual(context, {
+    start: true,
+    end: true,
+    asyncStart: 42,
+    asyncEnd: 42
+  });
+})();