webpack/components/WithOrganization/withOrganization.js in katello-3.8.1 vs webpack/components/WithOrganization/withOrganization.js in katello-3.9.0.rc1

@@ -1,28 +1,60 @@
 import React, { Component } from 'react';
-import { orgId } from '../../services/api';
+import PropTypes from 'prop-types';
+import { translate as __ } from 'foremanReact/common/I18n';
+import { get } from 'lodash';
 import SetOrganization from '../SelectOrg/SetOrganization';
-import titleWithCaret from '../../helpers/caret';
+import Header from '../../containers/Application/Headers';
 
 function withOrganization(WrappedComponent, redirectPath) {
-  return class CheckOrg extends Component {
+  class CheckOrg extends Component {
+    constructor(props) {
+      super(props);
+      this.state = { orgId: null };
+    }
+    static getDerivedStateFromProps(newProps, state) {
+      const orgNodeId = document.getElementById('organization-id').dataset.id;
+
+      if (state.orgId !== orgNodeId) {
+        return { orgId: orgNodeId };
+      }
+      return null;
+    }
+
     componentDidUpdate(prevProps) {
       const { location } = this.props;
+      const orgTitle = get(location, 'state.orgChanged');
+      const prevOrgTitle = get(prevProps, 'location.state.orgChanged');
 
-      // TODO: use topbar react component
-      const orgTitle = location.state && location.state.orgChanged;
-      const prevOrgTitle = prevProps.location.state && prevProps.location.state.orgChanged;
-
       if (orgTitle !== prevOrgTitle) {
-        document.getElementById('organization-dropdown').children[0].innerHTML = titleWithCaret(orgTitle);
+        window.tfm.nav.changeOrganization(orgTitle);
       }
     }
+
     render() {
-      if (!orgId()) {
-        return <SetOrganization redirectPath={redirectPath} />;
+      const { location } = this.props;
+      const newOrgSelected = get(location, 'state.orgChanged');
+
+      if (newOrgSelected) {
+        return <WrappedComponent {...this.props} />;
+      } else if (this.state.orgId === '') {
+        return (
+          <React.Fragment>
+            <Header title={__('Select Organization')} />
+            <SetOrganization redirectPath={redirectPath} />
+          </React.Fragment>);
       }
       return <WrappedComponent {...this.props} />;
     }
+  }
+
+  CheckOrg.propTypes = {
+    location: PropTypes.shape({}),
   };
+
+  CheckOrg.defaultProps = {
+    location: undefined,
+  };
+  return CheckOrg;
 }
 
 export default withOrganization;