feat: Fixed code after flatenning segment tree

tests/integration/config-capture-scalars.test.js

@@ -113,7 +113,7 @@ tests.push({
       }
 
       // Exact match to ensure no extra fields snuck in
-      assertSegments(transaction.trace.root, expectedSegments, { exact: true })
+      assertSegments(transaction.trace, transaction.trace.root, expectedSegments, { exact: true })
     })
 
     executeQuery(serverUrl, query, (err) => {

tests/lib/custom-assertions.js

@@ -63,6 +63,64 @@ function assertMetrics(
 }
 
 /**
+ * This function is used to verify that a tree of trace segments matches an
+ * expected tree of segment names. For example, if the trace looks like (i.e
+ * the `parent` parameter):
+ *
+ * ```js
+ * {
+ *   name: 'root-segment',
+ *   children: [
+ *    {
+ *      name: 'child 1',
+ *      children: [
+ *        {
+ *          name: 'grandchild 1',
+ *          children: [
+ *            {
+ *              name: 'great-grandchild',
+ *              children: []
+ *            }
+ *          ]
+ *        },
+ *        {
+ *          name: 'grandchild 2',
+ *          children: []
+ *        }
+ *      ]
+ *    },
+ *    {
+ *      name: 'child 2',
+ *      children: []
+ *    }
+ *   ]
+ * }
+ * ```
+ *
+ * Then the provided `expected` parameter should look like:
+ *
+ * ```js
+ * [
+ *   'root-segment',
+ *   [
+ *     'child 1',
+ *     [
+ *      'grandchild 1',
+ *      ['great-grandchild],
+ *      'grandchild 2'
+ *     ],
+ *     'child 2'
+ *   ],
+ * ]
+ * ```
+ *
+ * Ordering of the elements in the `expected` parameter is significant when
+ * `options.exact = true`. Regardless of the `exact` value, ordering of elements
+ * is significant to indicate the nesting order. Any string immediately
+ * followed by an array of strings indicates that the first string is a parent
+ * element, and the subsequent array of strings is its child elements.
+ *
+ * @param {Trace} trace             Transaction trace
  * @param {TraceSegment} parent     Parent segment
  * @param {Array} expected          Array of strings that represent segment names.
  *                                  If an item in the array is another array, it
@@ -78,7 +136,13 @@ function assertMetrics(
  * @param {object} [deps] Injected dependencies.
  * @param {object} [deps.assert] Assertion library to use.
  */
-function assertSegments(parent, expected, options, { assert = require('node:assert') } = {}) {
+function assertSegments(
+  trace,
+  parent,
+  expected,
+  options,
+  { assert = require('node:assert') } = {}
+) {
   let child
   let childCount = 0
 
@@ -91,7 +155,8 @@ function assertSegments(parent, expected, options, { assert = require('node:asse
   }
 
   function getChildren(_parent) {
-    return _parent.children.filter(function (item) {
+    const children = trace.getChildren(_parent.id)
+    return children.filter(function (item) {
       if (exact && options && options.exclude) {
         return options.exclude.indexOf(item.name) === -1
       }
@@ -126,7 +191,7 @@ function assertSegments(parent, expected, options, { assert = require('node:asse
           )
         }
       } else if (typeof sequenceItem === 'object') {
-        assertSegments(child, sequenceItem, options, { assert })
+        assertSegments(trace, child, sequenceItem, options, { assert })
       }
     }
 
@@ -138,14 +203,14 @@ function assertSegments(parent, expected, options, { assert = require('node:asse
 
       if (typeof sequenceItem === 'string') {
         // find corresponding child in parent
-        for (let j = 0; j < parent.children.length; j++) {
-          if (parent.children[j].name.startsWith(sequenceItem) === true) {
-            child = parent.children[j]
+        for (let j = 0; j < children.length; j++) {
+          if (children[j].name.startsWith(sequenceItem) === true) {
+            child = children[j]
           }
         }
         assert.ok(child, 'segment "' + parent.name + '" should have child "' + sequenceItem + '"')
         if (typeof expected[i + 1] === 'object') {
-          assertSegments(child, expected[i + 1], { exact }, { assert })
+          assertSegments(trace, child, expected[i + 1], { exact }, { assert })
         }
       }
     }

tests/lib/find-segment.js

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+module.exports = function findSegmentByName(trace, root, name) {
+  const children = trace.getChildren(root.id)
+  if (root.name === name) {
+    return root
+  } else if (children.length) {
+    for (let i = 0; i < children.length; i++) {
+      const child = children[i]
+      const found = findSegmentByName(trace, child, name)
+      if (found) {
+        return found
+      }
+    }
+  }
+  return null
+}

tests/unit/create-plugin.test.js

@@ -17,8 +17,8 @@ test('createPlugin edge cases', async (t) => {
     ctx.nr.operationSegment = {
       start: sinon.stub(),
       addAttribute: sinon.stub(),
-      transaction: { nameState: { setName: sinon.stub() } },
-      end: sinon.stub()
+      end: sinon.stub(),
+      start: sinon.stub()
     }
 
     ctx.nr.instrumentationApi = {

tests/versioned/apollo-federation/federated-gateway-server-setup.js

@@ -5,7 +5,6 @@
 
 'use strict'
 const utils = require('@newrelic/test-utilities')
-
 const federatedData = require('./federated-data-definitions')
 const { unloadModules } = require('../../lib/test-tools')
 

tests/versioned/apollo-federation/query-obfuscation.test.js

@@ -9,6 +9,7 @@ const assert = require('node:assert')
 const { setupFederatedGateway, teardownGateway } = require('./federated-gateway-server-setup')
 const { executeQuery } = require('../../lib/test-client')
 const { checkResult } = require('../common')
+const findSegmentByName = require('../../lib/find-segment')
 const SEGMENT_DESTINATION = 0x20
 const ANON_PLACEHOLDER = '<anonymous>'
 const QUERY_ATTRIBUTE_NAME = 'graphql.operation.query'
@@ -45,7 +46,7 @@ test('apollo-federation: query obfuscation', async (t) => {
     executeQuery(serverUrl, query, (err, result) => {
       assert.ok(!err)
       const operationName = `${OPERATION_PREFIX}/${ANON_PLACEHOLDER}/${path}`
-      const operationSegment = findSegmentByName(tx.trace.root, operationName)
+      const operationSegment = findSegmentByName(tx.trace, tx.trace.root, operationName)
 
       // only test one operation segment of three federated server transactions
       if (operationSegment) {
@@ -59,19 +60,3 @@ test('apollo-federation: query obfuscation', async (t) => {
     })
   })
 })
-
-function findSegmentByName(root, name) {
-  if (root.name === name) {
-    return root
-  } else if (root.children && root.children.length) {
-    for (let i = 0; i < root.children.length; i++) {
-      const child = root.children[i]
-      const found = findSegmentByName(child, name)
-      if (found) {
-        return found
-      }
-    }
-  }
-
-  return null
-}

tests/versioned/apollo-federation/segments.test.js

@@ -72,7 +72,7 @@ test('apollo-federation: federated segments', async (t) => {
         ]
       ]
 
-      assertSegments(tx.trace.root, expectedSegments, { exact: false }, { assert: plan })
+      assertSegments(tx.trace, tx.trace.root, expectedSegments, { exact: false }, { assert: plan })
       checkResult(plan, result, () => {})
     })
     await plan.completed
@@ -141,7 +141,7 @@ test('apollo-federation: federated segments', async (t) => {
         ]
       ]
 
-      assertSegments(tx.trace.root, expectedSegments, { exact: false }, { assert: plan })
+      assertSegments(tx.trace, tx.trace.root, expectedSegments, { exact: false }, { assert: plan })
       checkResult(plan, result, () => {
         plan.equal(result.length, 2)
       })