From 915f16fd030de0ea324c3abdc98ccc25b67d09fc Mon Sep 17 00:00:00 2001
From: Mark Berry <markberry867@gmail.com>
Date: Wed, 18 Nov 2020 00:56:42 -0600
Subject: [PATCH] Re-emit story data in ngZone

---
 .../angular/components/app.component.ts       |  4 +--
 .../src/client/preview/angular/helpers.ts     | 27 +++++++++++++++----
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/app/angular/src/client/preview/angular/components/app.component.ts b/app/angular/src/client/preview/angular/components/app.component.ts
index d70527b0767..57bd79442ab 100644
--- a/app/angular/src/client/preview/angular/components/app.component.ts
+++ b/app/angular/src/client/preview/angular/components/app.component.ts
@@ -15,7 +15,6 @@ import {
   SimpleChanges,
   SimpleChange,
   ChangeDetectorRef,
-  NgZone,
 } from '@angular/core';
 import { Observable, Subscription } from 'rxjs';
 import { first } from 'rxjs/operators';
@@ -39,7 +38,6 @@ export class AppComponent implements OnInit, OnDestroy {
   constructor(
     private cfr: ComponentFactoryResolver,
     private changeDetectorRef: ChangeDetectorRef,
-    private ngZone: NgZone,
     @Inject(STORY) private data: Observable<StoryFnAngularReturnType>
   ) {}
 
@@ -55,7 +53,7 @@ export class AppComponent implements OnInit, OnDestroy {
       );
 
       this.subscription = this.data.subscribe((newData) => {
-        this.ngZone.run(() => this.setProps(instance, newData));
+        this.setProps(instance, newData);
         childChangeDetectorRef.markForCheck();
         // Must detect changes on the current component in order to update any changes in child component's @HostBinding properties (angular/angular#22560)
         this.changeDetectorRef.detectChanges();
diff --git a/app/angular/src/client/preview/angular/helpers.ts b/app/angular/src/client/preview/angular/helpers.ts
index fa39fdcc542..97d41f65bfc 100644
--- a/app/angular/src/client/preview/angular/helpers.ts
+++ b/app/angular/src/client/preview/angular/helpers.ts
@@ -1,10 +1,10 @@
 // @ts-ignore
 import { document } from 'global';
-import { enableProdMode, NgModule, Component, NgModuleRef, Type } from '@angular/core';
+import { enableProdMode, NgModule, Component, NgModuleRef, Type, NgZone } from '@angular/core';
 import { FormsModule } from '@angular/forms';
 import { platformBrowserDynamic } from '@angular/platform-browser-dynamic';
 import { BrowserModule } from '@angular/platform-browser';
-import { ReplaySubject } from 'rxjs';
+import { Observable, ReplaySubject, Subscriber } from 'rxjs';
 import { StoryFn } from '@storybook/addons';
 import { AppComponent } from './components/app.component';
 import { STORY } from './app.token';
@@ -18,13 +18,27 @@ declare global {
 
 let platform: any = null;
 let promises: Promise<NgModuleRef<any>>[] = [];
-let storyData = new ReplaySubject(1);
+let storyData = new ReplaySubject<StoryFnAngularReturnType>(1);
 
 const moduleClass = class DynamicModule {};
 const componentClass = class DynamicComponent {};
 
 type DynamicComponentType = typeof componentClass;
 
+function storyDataFactory<T>(storyData: Observable<T>) {
+  return (ngZone: NgZone) => new Observable((subscriber: Subscriber<T>) => {
+    const sub = storyData.subscribe(
+      (v: T) => { ngZone.run(() => subscriber.next(v)); },
+      (err) => { ngZone.run(() => subscriber.error(err)); },
+      () => { ngZone.run(() => subscriber.complete()); }
+    );
+
+    return () => {
+      sub.unsubscribe();
+    };
+  });
+};
+
 const getModule = (
   declarations: (Type<any> | any[])[],
   entryComponents: (Type<any> | any[])[],
@@ -34,13 +48,16 @@ const getModule = (
 ) => {
   // Complete last ReplaySubject and create a new one for the current module
   storyData.complete();
-  storyData = new ReplaySubject(1);
+  storyData = new ReplaySubject<StoryFnAngularReturnType>(1);
   storyData.next(data);
 
   const moduleMeta = {
     declarations: [...declarations, ...(moduleMetadata.declarations || [])],
     imports: [BrowserModule, FormsModule, ...(moduleMetadata.imports || [])],
-    providers: [{ provide: STORY, useValue: storyData }, ...(moduleMetadata.providers || [])],
+    providers: [
+      { provide: STORY, useFactory: storyDataFactory(storyData.asObservable()), deps: [ NgZone ] },
+      ...(moduleMetadata.providers || [])
+    ],
     entryComponents: [...entryComponents, ...(moduleMetadata.entryComponents || [])],
     schemas: [...(moduleMetadata.schemas || [])],
     bootstrap: [...bootstrap],