fix(ivy): avoid using stale cache in TestBed if module overrides are defined (#33787)

NgModule compilation in JIT mode (that is also used in TestBed) caches module scopes on NgModule defs (using `transitiveCompileScopes` field). Module overrides (defined via TestBed.overrideModule) may invalidate this data by adding/removing items in `declarations` list. This commit forces TestBed to recalculate transitive scopes in case module overrides are present, so TestBed always gets the most up-to-date information.

PR Close #33787

packages/core/src/render3/jit/module.ts

@@ -420,19 +420,25 @@ export function patchComponentDefWithScope<C>(
 /**
  * Compute the pair of transitive scopes (compilation scope and exported scope) for a given module.
  *
- * This operation is memoized and the result is cached on the module's definition. It can be called
+ * By default this operation is memoized and the result is cached on the module's definition. You
- * on modules with components that have not fully compiled yet, but the result should not be used
+ * can avoid memoization and previously stored results (if available) by providing the second
- * until they have.
+ * argument with the `true` value (forcing transitive scopes recalculation).
+ *
+ * This function can be called on modules with components that have not fully compiled yet, but the
+ * result should not be used until they have.
+ *
+ * @param moduleType module that transitive scope should be calculated for.
+ * @param forceRecalc flag that indicates whether previously calculated and memoized values should
+ * be ignored and transitive scope to be fully recalculated.
  */
 export function transitiveScopesFor<T>(
-    moduleType: Type<T>,
+    moduleType: Type<T>, forceRecalc: boolean = false): NgModuleTransitiveScopes {
-    processNgModuleFn?: (ngModule: NgModuleType) => void): NgModuleTransitiveScopes {
   if (!isNgModule(moduleType)) {
     throw new Error(`${moduleType.name} does not have a module def (ɵmod property)`);
   }
   const def = getNgModuleDef(moduleType) !;
 
-  if (def.transitiveCompileScopes !== null) {
+  if (!forceRecalc && def.transitiveCompileScopes !== null) {
     return def.transitiveCompileScopes;
   }
 
@@ -471,13 +477,9 @@ export function transitiveScopesFor<T>(
       throw new Error(`Importing ${importedType.name} which does not have a ɵmod property`);
     }
 
-    if (processNgModuleFn) {
-      processNgModuleFn(importedType as NgModuleType);
-    }
-
     // When this module imports another, the imported module's exported directives and pipes are
     // added to the compilation scope of this module.
-    const importedScope = transitiveScopesFor(importedType, processNgModuleFn);
+    const importedScope = transitiveScopesFor(importedType, forceRecalc);
     importedScope.exported.directives.forEach(entry => scopes.compilation.directives.add(entry));
     importedScope.exported.pipes.forEach(entry => scopes.compilation.pipes.add(entry));
   });
@@ -496,7 +498,7 @@ export function transitiveScopesFor<T>(
     if (isNgModule(exportedType)) {
       // When this module exports another, the exported module's exported directives and pipes are
       // added to both the compilation and exported scopes of this module.
-      const exportedScope = transitiveScopesFor(exportedType, processNgModuleFn);
+      const exportedScope = transitiveScopesFor(exportedType, forceRecalc);
       exportedScope.exported.directives.forEach(entry => {
         scopes.compilation.directives.add(entry);
         scopes.exported.directives.add(entry);
@@ -512,7 +514,9 @@ export function transitiveScopesFor<T>(
     }
   });
 
+  if (!forceRecalc) {
     def.transitiveCompileScopes = scopes;
+  }
   return scopes;
 }
 

packages/core/test/test_bed_spec.ts

@@ -6,7 +6,7 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {Compiler, Component, Directive, ErrorHandler, Inject, Injectable, InjectionToken, Injector, Input, ModuleWithProviders, NgModule, Optional, Pipe, ɵsetClassMetadata as setClassMetadata, ɵɵdefineComponent as defineComponent, ɵɵdefineNgModule as defineNgModule, ɵɵtext as text} from '@angular/core';
+import {Compiler, Component, Directive, ErrorHandler, Inject, Injectable, InjectionToken, Injector, Input, ModuleWithProviders, NgModule, Optional, Pipe, ViewChild, ɵsetClassMetadata as setClassMetadata, ɵɵdefineComponent as defineComponent, ɵɵdefineNgModule as defineNgModule, ɵɵtext as text} from '@angular/core';
 import {TestBed, getTestBed} from '@angular/core/testing/src/test_bed';
 import {By} from '@angular/platform-browser';
 import {expect} from '@angular/platform-browser/testing/src/matchers';
@@ -289,6 +289,59 @@ describe('TestBed', () => {
     expect(SimpleService.ngOnDestroyCalls).toBe(0);
   });
 
+  describe('module overrides using TestBed.overrideModule', () => {
+    @Component({
+      selector: 'test-cmp',
+      template: '...',
+    })
+    class TestComponent {
+      testField = 'default';
+    }
+
+    @NgModule({
+      declarations: [TestComponent],
+      exports: [TestComponent],
+    })
+    class TestModule {
+    }
+
+    @Component({
+      selector: 'app-root',
+      template: `<test-cmp #testCmpCtrl></test-cmp>`,
+    })
+    class AppComponent {
+      @ViewChild('testCmpCtrl', {static: true}) testCmpCtrl !: TestComponent;
+    }
+
+    @NgModule({
+      declarations: [AppComponent],
+      imports: [TestModule],
+    })
+    class AppModule {
+    }
+    @Component({
+      selector: 'test-cmp',
+      template: '...',
+    })
+    class MockTestComponent {
+      testField = 'overwritten';
+    }
+
+    it('should allow declarations override', () => {
+      TestBed.configureTestingModule({
+        imports: [AppModule],
+      });
+      // replace TestComponent with MockTestComponent
+      TestBed.overrideModule(TestModule, {
+        remove: {declarations: [TestComponent], exports: [TestComponent]},
+        add: {declarations: [MockTestComponent], exports: [MockTestComponent]}
+      });
+      const fixture = TestBed.createComponent(AppComponent);
+      const app = fixture.componentInstance;
+      expect(app.testCmpCtrl.testField).toBe('overwritten');
+    });
+  });
+
   describe('multi providers', () => {
     const multiToken = new InjectionToken<string[]>('multiToken');
     const singleToken = new InjectionToken<string>('singleToken');

packages/core/testing/src/r3_test_bed_compiler.ts

@@ -88,6 +88,7 @@ export class R3TestBedCompiler {
 
   private testModuleType: NgModuleType<any>;
   private testModuleRef: NgModuleRef<any>|null = null;
+  private hasModuleOverrides: boolean = false;
 
   constructor(private platform: PlatformRef, private additionalModuleTypes: Type<any>|Type<any>[]) {
     class DynamicTestModule {}
@@ -122,6 +123,8 @@ export class R3TestBedCompiler {
   }
 
   overrideModule(ngModule: Type<any>, override: MetadataOverride<NgModule>): void {
+    this.hasModuleOverrides = true;
+
     // Compile the module right away.
     this.resolvers.module.addOverride(ngModule, override);
     const metadata = this.resolvers.module.resolve(ngModule);
@@ -331,8 +334,15 @@ export class R3TestBedCompiler {
     const getScopeOfModule =
         (moduleType: Type<any>| TestingModuleOverride): NgModuleTransitiveScopes => {
           if (!moduleToScope.has(moduleType)) {
-            const realType = isTestingModuleOverride(moduleType) ? this.testModuleType : moduleType;
+            const isTestingModule = isTestingModuleOverride(moduleType);
-            moduleToScope.set(moduleType, transitiveScopesFor(realType));
+            const realType = isTestingModule ? this.testModuleType : moduleType as Type<any>;
+            // Module overrides (via TestBed.overrideModule) might affect scopes that were
+            // previously calculated and stored in `transitiveCompileScopes`. If module overrides
+            // are present, always re-calculate transitive scopes to have the most up-to-date
+            // information available. The `moduleToScope` map avoids repeated re-calculation of
+            // scopes for the same module.
+            const forceRecalc = !isTestingModule && this.hasModuleOverrides;
+            moduleToScope.set(moduleType, transitiveScopesFor(realType, forceRecalc));
           }
           return moduleToScope.get(moduleType) !;
         };