[test-optimization] [SDTEST-1060] Add mocked files of Jest to the list of dependencies in TIA (#5885)

Author: Mariovido

integration-tests/ci-visibility/jest/mocked-test.js

@@ -0,0 +1,6 @@
+// eslint-disable-next-line no-undef
+jest.mock('../test/sum.js')
+
+test('adds 1 + 2 to equal 3', () => {
+  expect(1 + 2).toBe(3)
+})

integration-tests/jest/jest.spec.js

@@ -1560,6 +1560,37 @@ describe('jest CommonJS', () => {
         }).catch(done)
       })
     })
+
+    it('report code coverage with all mocked files', (done) => {
+      const codeCovRequestPromise = receiver.payloadReceived(({ url }) => url === '/api/v2/citestcov')
+
+      codeCovRequestPromise.then((codeCovRequest) => {
+        const allCoverageFiles = codeCovRequest.payload
+          .flatMap(coverage => coverage.content.coverages)
+          .flatMap(file => file.files)
+          .map(file => file.filename)
+
+        assert.includeMembers(allCoverageFiles, [
+          'ci-visibility/test/sum.js',
+          'ci-visibility/jest/mocked-test.js'
+        ])
+      }).catch(done)
+
+      childProcess = exec(
+        runTestsWithCoverageCommand,
+        {
+          cwd,
+          env: {
+            ...getCiVisAgentlessConfig(receiver.port),
+            TESTS_TO_RUN: 'jest/mocked-test.js',
+          },
+          stdio: 'pipe'
+        }
+      )
+      childProcess.on('exit', () => {
+        done()
+      })
+    })
   })
 
   context('early flake detection', () => {

packages/datadog-instrumentations/src/jest.js

@@ -3,6 +3,7 @@
 const { addHook, channel } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
+const path = require('path')
 const {
   getCoveredFilenamesFromCoverage,
   JEST_WORKER_TRACE_PAYLOAD_CODE,
@@ -90,6 +91,7 @@ const retriedTestsToNumAttempts = new Map()
 const newTestsTestStatuses = new Map()
 const attemptToFixRetriedTestsStatuses = new Map()
 const wrappedWorkers = new WeakSet()
+const testSuiteMockedFiles = new Map()
 
 const BREAKPOINT_HIT_GRACE_PERIOD_MS = 200
 
@@ -137,6 +139,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       this.nameToParams = {}
       this.global._ddtrace = global._ddtrace
       this.hasSnapshotTests = undefined
+      this.testSuiteAbsolutePath = context.testPath
 
       this.displayName = config.projectConfig?.displayName?.name
       this.testEnvironmentOptions = getTestEnvironmentOptions(config)
@@ -1095,10 +1098,12 @@ function jestAdapterWrapper (jestAdapter, jestVersion) {
       if (environment.testEnvironmentOptions?._ddTestCodeCoverageEnabled) {
         const root = environment.repositoryRoot || environment.rootDir
 
-        const coverageFiles = getCoveredFilenamesFromCoverage(environment.global.__coverage__)
-          .map(filename => getTestSuitePath(filename, root))
+        const getFilesWithPath = (files) => files.map(file => getTestSuitePath(file, root))
 
-        testSuiteCodeCoverageCh.publish({ coverageFiles, testSuite: environment.testSourceFile })
+        const coverageFiles = getFilesWithPath(getCoveredFilenamesFromCoverage(environment.global.__coverage__))
+        const mockedFiles = getFilesWithPath(testSuiteMockedFiles.get(environment.testSuiteAbsolutePath) || [])
+
+        testSuiteCodeCoverageCh.publish({ coverageFiles, testSuite: environment.testSourceFile, mockedFiles })
       }
       testSuiteFinishCh.publish({ status, errorMessage })
       return suiteResults
@@ -1267,6 +1272,23 @@ addHook({
 }, (runtimePackage) => {
   const Runtime = runtimePackage.default ?? runtimePackage
 
+  shimmer.wrap(Runtime.prototype, '_createJestObjectFor', _createJestObjectFor => function (from) {
+    const result = _createJestObjectFor.apply(this, arguments)
+    const suiteFilePath = this._testPath
+
+    shimmer.wrap(result, 'mock', mock => function (moduleName) {
+      if (suiteFilePath) {
+        const existingMockedFiles = testSuiteMockedFiles.get(suiteFilePath) || []
+        const suiteDir = path.dirname(suiteFilePath)
+        const mockPath = path.resolve(suiteDir, moduleName)
+        existingMockedFiles.push(mockPath)
+        testSuiteMockedFiles.set(suiteFilePath, existingMockedFiles)
+      }
+      return mock.apply(this, arguments)
+    })
+    return result
+  })
+
   shimmer.wrap(Runtime.prototype, 'requireModuleOrMock', requireModuleOrMock => function (from, moduleName) {
     // TODO: do this for every library that we instrument
     if (shouldBypassJestRequireEngine(moduleName)) {

packages/datadog-plugin-jest/src/index.js

@@ -318,11 +318,11 @@ class JestPlugin extends CiPlugin {
      * because this subscription happens in a different process from the one
      * fetching the ITR config.
      */
-    this.addSub('ci:jest:test-suite:code-coverage', ({ coverageFiles, testSuite }) => {
+    this.addSub('ci:jest:test-suite:code-coverage', ({ coverageFiles, testSuite, mockedFiles }) => {
       if (!coverageFiles.length) {
         this.telemetry.count(TELEMETRY_CODE_COVERAGE_EMPTY)
       }
-      const files = [...coverageFiles, testSuite]
+      const files = [...coverageFiles, ...mockedFiles, testSuite]
 
       const { _traceId, _spanId } = this.testSuiteSpan.context()
       const formattedCoverage = {