fix(ivy): Only restore registered modules if user compiles modules with TestBed (#32944)

There are a couple scenarios that are problematic and need special
handling:

1. A user has a custom implementation of lazy-loaded modules, sets some
provider overrides, then compiles the module so it can be loaded. In a
follow-up test, the user sets different overrides for the module and
then compiles. This is problematic because we need to be sure the module
registered in the first test is not used, so we need to clear it out of
the modules list in `ng_module_factory_registration`.
2. A user has a similar lazy-loaded module factory implementation but
relies on the module being registered automatically. This can happen,
for example, as a side effect of importing the ngfactory file.

PR Close #32944

packages/core/src/linker/ng_module_factory_registration.ts

@@ -12,13 +12,14 @@ import {stringify} from '../util/stringify';
 
 import {NgModuleFactory} from './ng_module_factory';
 
+export type ModuleRegistrationMap = Map<string, NgModuleFactory<any>|NgModuleType>;
 
 /**
  * Map of module-id to the corresponding NgModule.
  * - In pre Ivy we track NgModuleFactory,
  * - In post Ivy we track the NgModuleType
  */
-const modules = new Map<string, NgModuleFactory<any>|NgModuleType>();
+let modules: ModuleRegistrationMap = new Map();
 
 /**
  * Registers a loaded module. Should only be called from generated NgModuleFactory code.
@@ -54,10 +55,18 @@ export function registerNgModuleType(ngModuleType: NgModuleType) {
   }
 }
 
-export function clearModuleRegistry(): void {
+export function clearRegisteredModuleState(): void {
   modules.clear();
 }
 
+export function getRegisteredModulesState(): ModuleRegistrationMap {
+  return new Map(modules);
+}
+
+export function restoreRegisteredModulesState(moduleMap: ModuleRegistrationMap) {
+  modules = new Map(moduleMap);
+}
+
 export function getRegisteredNgModuleType(id: string) {
   return modules.get(id);
 }

packages/core/test/linker/ng_module_integration_spec.ts

@@ -17,7 +17,7 @@ import {expect} from '@angular/platform-browser/testing/src/matchers';
 import {modifiedInIvy, obsoleteInIvy, onlyInIvy} from '@angular/private/testing';
 
 import {InternalNgModuleRef, NgModuleFactory} from '../../src/linker/ng_module_factory';
-import {clearModuleRegistry} from '../../src/linker/ng_module_factory_registration';
+import {clearRegisteredModuleState} from '../../src/linker/ng_module_factory_registration';
 import {stringify} from '../../src/util/stringify';
 
 class Engine {}
@@ -294,11 +294,7 @@ function declareTests(config?: {useJit: boolean}) {
     describe('id', () => {
       const token = 'myid';
 
-      // Ivy TestBed clears module registry in resetTestingModule so this afterEach is not needed
-      // for Ivy
-      if (!ivyEnabled) {
-        afterEach(() => clearModuleRegistry());
-      }
+      afterEach(() => clearRegisteredModuleState());
 
       it('should register loaded modules', () => {
         @NgModule({id: token})

packages/core/test/test_bed_spec.ts

@@ -7,6 +7,8 @@
  */
 
 import {Compiler, Component, Directive, ErrorHandler, Inject, Injectable, InjectionToken, Input, ModuleWithProviders, NgModule, Optional, Pipe, getModuleFactory, ɵsetClassMetadata as setClassMetadata, ɵɵdefineComponent as defineComponent, ɵɵdefineNgModule as defineNgModule, ɵɵtext as text} from '@angular/core';
+import {registerModuleFactory} from '@angular/core/src/linker/ng_module_factory_registration';
+import {NgModuleFactory} from '@angular/core/src/render3';
 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';
@@ -732,15 +734,30 @@ describe('TestBed', () => {
       });
 
   onlyInIvy('Ivy module registration happens when NgModuleFactory is created')
-      .it('cleans up registered modules', async() => {
-        @NgModule({id: 'my_module'})
-        class MyModule {
-        }
+      .describe('cleans up registered modules - ', () => {
+        it('removes modules registered with TestBed', async() => {
+          @NgModule({id: 'my_module'})
+          class MyModule {
+          }
 
-        expect(() => getModuleFactory('my_module')).toThrowError();
-        await TestBed.inject(Compiler).compileModuleAsync(MyModule);
-        expect(() => getModuleFactory('my_module')).not.toThrowError();
-        TestBed.resetTestingModule();
-        expect(() => getModuleFactory('my_module')).toThrowError();
+          expect(() => getModuleFactory('my_module')).toThrowError();
+          await TestBed.inject(Compiler).compileModuleAsync(MyModule);
+          expect(() => getModuleFactory('my_module')).not.toThrowError();
+          TestBed.resetTestingModule();
+          expect(() => getModuleFactory('my_module')).toThrowError();
+        });
+
+        it('does not remove modules registered outside TestBed (i.e., side effect registration in ngfactory files)',
+           () => {
+             @NgModule({id: 'auto_module'})
+             class AutoModule {
+             }
+
+             expect(() => getModuleFactory('auto_module')).toThrowError();
+             registerModuleFactory('auto_module', new NgModuleFactory(AutoModule));
+             expect(() => getModuleFactory('auto_module')).not.toThrowError();
+             TestBed.resetTestingModule();
+             expect(() => getModuleFactory('auto_module')).not.toThrowError();
+           });
       });
 });

packages/core/testing/src/r3_test_bed.ts

@@ -35,7 +35,7 @@ import {MetadataOverride} from './metadata_override';
 import {TestBed} from './test_bed';
 import {ComponentFixtureAutoDetect, ComponentFixtureNoNgZone, TestBedStatic, TestComponentRenderer, TestModuleMetadata} from './test_bed_common';
 import {R3TestBedCompiler} from './r3_test_bed_compiler';
-import {clearModuleRegistry} from '../../src/linker/ng_module_factory_registration';
+import {clearRegisteredModuleState} from '../../src/linker/ng_module_factory_registration';
 
 let _nextRootElementId = 0;
 
@@ -230,7 +230,6 @@ export class TestBedRender3 implements TestBed {
   }
 
   resetTestingModule(): void {
-    clearModuleRegistry();
     this.checkGlobalCompilationFinished();
     resetCompiledComponents();
     if (this._compiler !== null) {

packages/core/testing/src/r3_test_bed_compiler.ts

@@ -8,6 +8,7 @@
 
 import {ResourceLoader} from '@angular/compiler';
 import {ApplicationInitStatus, COMPILER_OPTIONS, Compiler, Component, Directive, Injector, LOCALE_ID, ModuleWithComponentFactories, ModuleWithProviders, NgModule, NgModuleFactory, NgZone, Pipe, PlatformRef, Provider, Type, ɵDEFAULT_LOCALE_ID as DEFAULT_LOCALE_ID, ɵDirectiveDef as DirectiveDef, ɵNG_COMPONENT_DEF as NG_COMPONENT_DEF, ɵNG_DIRECTIVE_DEF as NG_DIRECTIVE_DEF, ɵNG_INJECTOR_DEF as NG_INJECTOR_DEF, ɵNG_MODULE_DEF as NG_MODULE_DEF, ɵNG_PIPE_DEF as NG_PIPE_DEF, ɵNgModuleFactory as R3NgModuleFactory, ɵNgModuleTransitiveScopes as NgModuleTransitiveScopes, ɵNgModuleType as NgModuleType, ɵRender3ComponentFactory as ComponentFactory, ɵRender3NgModuleRef as NgModuleRef, ɵcompileComponent as compileComponent, ɵcompileDirective as compileDirective, ɵcompileNgModuleDefs as compileNgModuleDefs, ɵcompilePipe as compilePipe, ɵgetInjectableDef as getInjectableDef, ɵpatchComponentDefWithScope as patchComponentDefWithScope, ɵsetLocaleId as setLocaleId, ɵtransitiveScopesFor as transitiveScopesFor, ɵɵInjectableDef as InjectableDef} from '@angular/core';
+import {ModuleRegistrationMap, getRegisteredModulesState, restoreRegisteredModulesState} from '../../src/linker/ng_module_factory_registration';
 
 import {clearResolutionOfComponentResourcesQueue, isComponentDefPendingResolution, resolveComponentResources, restoreComponentResolutionQueue} from '../../src/metadata/resource_loading';
 
@@ -41,6 +42,7 @@ interface CleanupOperation {
 
 export class R3TestBedCompiler {
   private originalComponentResolutionQueue: Map<Type<any>, Component>|null = null;
+  private originalRegisteredModules: null|ModuleRegistrationMap = null;
 
   // Testing module configuration
   private declarations: Type<any>[] = [];
@@ -264,6 +266,9 @@ export class R3TestBedCompiler {
    * @internal
    */
   async _compileNgModuleAsync(moduleType: Type<any>): Promise<void> {
+    if (this.originalRegisteredModules === null) {
+      this.originalRegisteredModules = getRegisteredModulesState();
+    }
     this.queueTypesFromModulesArray([moduleType]);
     await this.compileComponents();
     this.applyProviderOverrides();
@@ -535,6 +540,10 @@ export class R3TestBedCompiler {
     this.initialNgDefs.clear();
     this.moduleProvidersOverridden.clear();
     this.restoreComponentResolutionQueue();
+    if (this.originalRegisteredModules) {
+      restoreRegisteredModulesState(this.originalRegisteredModules);
+      this.originalRegisteredModules = null;
+    }
     // Restore the locale ID to the default value, this shouldn't be necessary but we never know
     setLocaleId(DEFAULT_LOCALE_ID);
   }