feat: withEnergy function component support

arwes · Oct 26, 2019 · d5c467e · d5c467e
1 parent ae192c0
commit d5c467e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 15 deletions.
diff --git a/packages/animation/src/withEnergy/withEnergy.js b/packages/animation/src/withEnergy/withEnergy.js
@@ -14,6 +14,23 @@ function withEnergy (options) {
     function EnergyManager ({ forwardRef, children, ...etc }) {
       const energy = useEnergy();
       const inner = useRef();
+      const others = {};
+
+      // Only if component requires a reference, add it.
+      // This allows the component to be a function if not.
+      if (cycles || forwardRef) {
+        others.ref = ref => {
+          inner.current = ref;
+          if (forwardRef) {
+            if (typeof forwardRef === 'function') {
+              forwardRef(ref);
+            }
+            else {
+              forwardRef.current = ref;
+            }
+          }
+        };
+      }
 
       useEffect(() => {
         if (inner && cycles) {
@@ -27,21 +44,7 @@ function withEnergy (options) {
       }, [energy.flow.value]);
 
       return (
-        <Inner
-          {...etc}
-          ref={ref => {
-            inner.current = ref;
-            if (forwardRef) {
-              if (typeof forwardRef === 'function') {
-                forwardRef(ref);
-              }
-              else {
-                forwardRef.current = ref;
-              }
-            }
-          }}
-          energy={energy}
-        >
+        <Inner {...etc} {...others} energy={energy}>
           {children}
         </Inner>
       );

diff --git a/packages/animation/src/withEnergy/withEnergy.test.js b/packages/animation/src/withEnergy/withEnergy.test.js
@@ -141,3 +141,10 @@ test('Should set component default props and be extended by provided props', ()
   const Example = withEnergy(options)(ExampleComponent);
   render(<Example energy={{ duration: { enter: 500 } }} />);
 });
+
+test('Should accept functions as components when "cycles=false" and not referenced', () => {
+  const ExampleComponent = () => <div />;
+  const options = { cycles: false };
+  const Example = withEnergy(options)(ExampleComponent);
+  render(<Example />);
+});