Fix flaky feature flag test.

tensorboard/webapp/BUILD

@@ -264,7 +264,6 @@ tf_ng_web_test_suite(
         "//tensorboard/webapp/core/views:test_lib",
         "//tensorboard/webapp/customization:customization_test_lib",
         "//tensorboard/webapp/deeplink:deeplink_test_lib",
-        "//tensorboard/webapp/feature_flag/views:views_test",
         "//tensorboard/webapp/header:test_lib",
         "//tensorboard/webapp/metrics:integration_test",
         "//tensorboard/webapp/metrics:test_lib",
@@ -288,7 +287,6 @@ tf_ng_web_test_suite(
         "//tensorboard/webapp/runs/views/runs_table:runs_table_test",
         "//tensorboard/webapp/tbdev_upload:test_lib",
         "//tensorboard/webapp/util:util_tests",
-        "//tensorboard/webapp/webapp_data_source:feature_flag_test_lib",
         "//tensorboard/webapp/webapp_data_source:http_client_test",
         "//tensorboard/webapp/webapp_data_source:webapp_data_source_test_lib",
         "//tensorboard/webapp/widgets:resize_detector_test",

tensorboard/webapp/feature_flag/BUILD

@@ -43,10 +43,9 @@ tf_ts_library(
 tf_ng_web_test_suite(
     name = "karma_test",
     deps = [
-        ":test_lib",
         "//tensorboard/webapp/feature_flag/effects:effects_test_lib",
         "//tensorboard/webapp/feature_flag/http:http_test_lib",
         "//tensorboard/webapp/feature_flag/store:store_test_lib",
-        "//tensorboard/webapp/feature_flag/views:views_test",
+        "//tensorboard/webapp/feature_flag/views:views_test_lib",
     ],
 )

tensorboard/webapp/feature_flag/views/BUILD

@@ -1,4 +1,4 @@
-load("//tensorboard/defs:defs.bzl", "tf_ng_module", "tf_sass_binary", "tf_ts_library")
+load("//tensorboard/defs:defs.bzl", "tf_ng_module", "tf_sass_binary")
 
 package(default_visibility = ["//tensorboard:internal"])
 
@@ -23,6 +23,7 @@ tf_ng_module(
     deps = [
         ":feature_flag_dialog",
         "//tensorboard/webapp:app_state",
+        "//tensorboard/webapp/angular:expect_angular_cdk_overlay",
         "//tensorboard/webapp/angular:expect_angular_material_dialog",
         "//tensorboard/webapp/feature_flag/actions",
         "//tensorboard/webapp/feature_flag/store",
@@ -61,8 +62,8 @@ tf_ng_module(
     ],
 )
 
-tf_ts_library(
-    name = "views_test",
+tf_ng_module(
+    name = "views_test_lib",
     testonly = True,
     srcs = [
         "feature_flag_dialog_test.ts",

tensorboard/webapp/feature_flag/views/feature_flag_modal_trigger_container.ts

@@ -12,9 +12,11 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-import {Component, OnInit} from '@angular/core';
+import {ComponentType} from '@angular/cdk/overlay';
+import {Component, OnDestroy, OnInit} from '@angular/core';
 import {MatDialog, MatDialogRef} from '@angular/material/dialog';
 import {Store} from '@ngrx/store';
+import {Subject, takeUntil} from 'rxjs';
 import {State} from '../../app_state';
 import {featureFlagOverridesReset} from '../actions/feature_flag_actions';
 import {getShowFlagsEnabled} from '../store/feature_flag_selectors';
@@ -31,38 +33,54 @@ const util = {
   template: ``,
   styles: [],
 })
-export class FeatureFlagModalTriggerContainer implements OnInit {
+export class FeatureFlagModalTriggerContainer implements OnInit, OnDestroy {
+  // Allow the dialog component type to be overriden for testing purposes.
+  featureFlagDialogType: ComponentType<any> = FeatureFlagDialogContainer;
+
   readonly showFeatureFlags$ = this.store.select(getShowFlagsEnabled);
   private featureFlagsDialog?: MatDialogRef<FeatureFlagDialogContainer>;
+  private ngUnsubscribe = new Subject<void>();
 
   constructor(
     private readonly store: Store<State>,
     private dialog: MatDialog
   ) {}
 
   ngOnInit() {
-    this.showFeatureFlags$.subscribe((showFeatureFlags: boolean) => {
-      if (showFeatureFlags) {
-        this.featureFlagsDialog = this.dialog.open(FeatureFlagDialogContainer);
-        this.featureFlagsDialog.afterClosed().subscribe(() => {
-          // Reset the 'showFlags' flag when the dialog is closed to prevent the
-          // dialog from appearing again after the page is refreshed.
-          this.store.dispatch(
-            featureFlagOverridesReset({
-              flags: ['showFlags'],
-            })
+    this.showFeatureFlags$
+      .pipe(takeUntil(this.ngUnsubscribe))
+      .subscribe((showFeatureFlags: boolean) => {
+        if (showFeatureFlags) {
+          this.featureFlagsDialog = this.dialog.open(
+            this.featureFlagDialogType
           );
-          // Reload the page so that the application restarts with stable
-          // feature flag values.
-          // Wait one tick before reloading the page so the 'showFlags'
-          // reset has a chance to be reflected in the URL before page reload.
-          setTimeout(() => {
-            util.reloadWindow();
-          }, 1);
-        });
-        return;
-      }
-    });
+          this.featureFlagsDialog
+            .afterClosed()
+            .pipe(takeUntil(this.ngUnsubscribe))
+            .subscribe(() => {
+              // Reset the 'showFlags' flag when the dialog is closed to prevent the
+              // dialog from appearing again after the page is refreshed.
+              this.store.dispatch(
+                featureFlagOverridesReset({
+                  flags: ['showFlags'],
+                })
+              );
+              // Reload the page so that the application restarts with stable
+              // feature flag values.
+              // Wait one tick before reloading the page so the 'showFlags'
+              // reset has a chance to be reflected in the URL before page reload.
+              setTimeout(() => {
+                util.reloadWindow();
+              }, 1);
+            });
+          return;
+        }
+      });
+  }
+
+  ngOnDestroy() {
+    this.ngUnsubscribe.next();
+    this.ngUnsubscribe.complete();
   }
 }
 

tensorboard/webapp/feature_flag/views/feature_flag_modal_trigger_container_test.ts

@@ -14,14 +14,9 @@ limitations under the License.
 ==============================================================================*/
 import {HarnessLoader} from '@angular/cdk/testing';
 import {TestbedHarnessEnvironment} from '@angular/cdk/testing/testbed';
-import {NO_ERRORS_SCHEMA} from '@angular/core';
+import {Component} from '@angular/core';
 import {ComponentFixture, TestBed} from '@angular/core/testing';
-import {
-  MatDialog,
-  MatDialogModule,
-  MatDialogRef,
-  MAT_DIALOG_DATA,
-} from '@angular/material/dialog';
+import {MatDialogModule} from '@angular/material/dialog';
 import {MatDialogHarness} from '@angular/material/dialog/testing';
 import {NoopAnimationsModule} from '@angular/platform-browser/animations';
 import {Store} from '@ngrx/store';
@@ -38,27 +33,24 @@ import {
   FeatureFlagModalTriggerContainer,
   TEST_ONLY,
 } from './feature_flag_modal_trigger_container';
-import {FeatureFlagDialogContainer} from './feature_flag_dialog_container';
+
+@Component({
+  selector: 'testable-feature-flag-dialog-container',
+  template: '<div>Test</div>',
+})
+class TestableFeatureFlagDialogContainer {}
 
 describe('feature_flag_modal_trigger_container', () => {
   let store: MockStore<State>;
-
   let fixture: ComponentFixture<FeatureFlagModalTriggerContainer>;
-  let loader: HarnessLoader;
   let rootLoader: HarnessLoader;
 
   beforeEach(async () => {
     await TestBed.configureTestingModule({
       imports: [MatDialogModule, NoopAnimationsModule],
-      declarations: [FeatureFlagDialogContainer],
-      providers: [
-        provideMockTbStore(),
-        {provide: MatDialogRef, useValue: MatDialog},
-      ],
-      schemas: [NO_ERRORS_SCHEMA],
+      providers: [provideMockTbStore()],
     }).compileComponents();
 
-    TestBed.overrideProvider(MAT_DIALOG_DATA, {useValue: {}});
     store = TestBed.inject<Store<State>>(Store) as MockStore<State>;
 
     spyOn(store, 'dispatch').and.stub();
@@ -73,7 +65,11 @@ describe('feature_flag_modal_trigger_container', () => {
 
   function createComponent() {
     fixture = TestBed.createComponent(FeatureFlagModalTriggerContainer);
-    loader = TestbedHarnessEnvironment.loader(fixture);
+    // Override the dialog component to something that does not require a
+    // working store. The dialog component sometimes completes initialization
+    // after the test is torn down and the selectors have been reset.
+    fixture.componentInstance.featureFlagDialogType =
+      TestableFeatureFlagDialogContainer;
     rootLoader = TestbedHarnessEnvironment.documentRootLoader(fixture);
   }