bugfixing & polishing

1 files changed, 19 insertions, 3 deletions

diff --git a/src/components/ui/PremiumFeatureContainer/index.js b/src/components/ui/PremiumFeatureContainer/index.js
index 3c1e0fac3..8d2746e22 100644
--- a/src/components/ui/PremiumFeatureContainer/index.js
+++ b/src/components/ui/PremiumFeatureContainer/index.js
@@ -10,6 +10,7 @@ import UserStore from '../../../stores/UserStore';
 
 import styles from './styles';
 import { gaEvent } from '../../../lib/analytics';
+import { FeatureStore } from '../../../features/utils/FeatureStore';
 
 const messages = defineMessages({
   action: {
@@ -22,7 +23,10 @@ const messages = defineMessages({
 class PremiumFeatureContainer extends Component {
   static propTypes = {
     classes: PropTypes.object.isRequired,
-    condition: PropTypes.bool,
+    condition: PropTypes.oneOfType([
+      PropTypes.bool,
+      PropTypes.func,
+    ]),
     gaEventInfo: PropTypes.shape({
       category: PropTypes.string.isRequired,
       event: PropTypes.string.isRequired,
@@ -31,7 +35,7 @@ class PremiumFeatureContainer extends Component {
   };
 
   static defaultProps = {
-    condition: true,
+    condition: null,
     gaEventInfo: null,
   };
 
@@ -51,7 +55,18 @@ class PremiumFeatureContainer extends Component {
 
     const { intl } = this.context;
 
-    return !stores.user.data.isPremium && !!condition ? (
+    let showWrapper = !!condition;
+
+    if (condition === null) {
+      showWrapper = !stores.user.data.isPremium;
+    } else if (typeof condition === 'function') {
+      showWrapper = condition({
+        isPremium: stores.user.data.isPremium,
+        features: stores.features.features,
+      });
+    }
+
+    return showWrapper ? (
       <div className={classes.container}>
         <div className={classes.titleContainer}>
           <p className={classes.title}>Premium Feature</p>
@@ -81,6 +96,7 @@ PremiumFeatureContainer.wrappedComponent.propTypes = {
   children: oneOrManyChildElements.isRequired,
   stores: PropTypes.shape({
     user: PropTypes.instanceOf(UserStore).isRequired,
+    features: PropTypes.instanceOf(FeatureStore).isRequired,
   }).isRequired,
   actions: PropTypes.shape({
     ui: PropTypes.shape({