[ML] Fix license subscription race condition.

x-pack/plugins/ml/public/application/app.tsx

@@ -80,11 +80,11 @@ export const renderApp = (
 
   deps.kibanaLegacy.loadFontAwesome();
 
-  const mlLicense = setLicenseCache(deps.licensing);
-
   appMountParams.onAppLeave((actions) => actions.default());
 
-  ReactDOM.render(<App coreStart={coreStart} deps={deps} />, appMountParams.element);
+  const mlLicense = setLicenseCache(deps.licensing, [
+    () => ReactDOM.render(<App coreStart={coreStart} deps={deps} />, appMountParams.element),
+  ]);
 
   return () => {
     mlLicense.unsubscribe();

x-pack/plugins/ml/public/application/license/check_license.tsx

@@ -5,6 +5,7 @@
  */
 
 import { LicensingPluginSetup } from '../../../../licensing/public';
+import { MlLicense } from '../../../common/license';
 import { MlClientLicense } from './ml_client_license';
 
 let mlLicense: MlClientLicense | null = null;
@@ -16,9 +17,12 @@ let mlLicense: MlClientLicense | null = null;
  * @param {LicensingPluginSetup} licensingSetup
  * @returns {MlClientLicense}
  */
-export function setLicenseCache(licensingSetup: LicensingPluginSetup) {
+export function setLicenseCache(
+  licensingSetup: LicensingPluginSetup,
+  postInitFunctions?: Array<(lic: MlLicense) => void>
+) {
   mlLicense = new MlClientLicense();
-  mlLicense.setup(licensingSetup.license$);
+  mlLicense.setup(licensingSetup.license$, postInitFunctions);
   return mlLicense;
 }
 

x-pack/plugins/ml/public/application/license/ml_client_license.test.ts

@@ -0,0 +1,59 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { Observable, Subject } from 'rxjs';
+import { ILicense } from '../../../../licensing/common/types';
+
+import { MlClientLicense } from './ml_client_license';
+
+describe('MlClientLicense', () => {
+  test('should miss the license update when initialized without postInitFunction', () => {
+    const mlLicense = new MlClientLicense();
+
+    // upon instantiation the full license doesn't get set
+    expect(mlLicense.isFullLicense()).toBe(false);
+
+    const license$ = new Subject();
+
+    mlLicense.setup(license$ as Observable<ILicense>);
+
+    // if the observable wasn't triggered the full license is still not set
+    expect(mlLicense.isFullLicense()).toBe(false);
+
+    license$.next({
+      check: () => ({ state: 'valid' }),
+      getFeature: () => ({ isEnabled: true }),
+      status: 'valid',
+    });
+
+    // once the observable triggered the license should be set
+    expect(mlLicense.isFullLicense()).toBe(true);
+  });
+
+  test('should not miss the license update when initialized with postInitFunction', (done) => {
+    const mlLicense = new MlClientLicense();
+
+    // upon instantiation the full license doesn't get set
+    expect(mlLicense.isFullLicense()).toBe(false);
+
+    const license$ = new Subject();
+
+    mlLicense.setup(license$ as Observable<ILicense>, [
+      (license) => {
+        // when passed in via postInitFunction callback, the license should be valid
+        // even if the license$ observable gets triggered after this setup.
+        expect(license.isFullLicense()).toBe(true);
+        done();
+      },
+    ]);
+
+    license$.next({
+      check: () => ({ state: 'valid' }),
+      getFeature: () => ({ isEnabled: true }),
+      status: 'valid',
+    });
+  });
+});