From 81534e1a3c5ccd0d8d522793ebccfbc9fa6271af Mon Sep 17 00:00:00 2001
From: Emmanuel Zamora <emmanuel.zamora@split.io>
Date: Mon, 14 Feb 2022 13:41:39 -0300
Subject: [PATCH] unit tests and fix memoize issue

---
 src/SplitClient.tsx                     |  1 +
 src/SplitTreatments.tsx                 |  7 ++-
 src/__tests__/SplitClient.test.tsx      | 76 +++++++++++++++++++++++++
 src/__tests__/testUtils/mockSplitSdk.ts | 12 ++++
 src/__tests__/useClient.test.tsx        | 75 ++++++++++++++++++++++++
 src/__tests__/withSplitClient.test.tsx  | 74 ++++++++++++++++++++++++
 src/useClient.ts                        |  6 +-
 7 files changed, 245 insertions(+), 6 deletions(-)

diff --git a/src/SplitClient.tsx b/src/SplitClient.tsx
index 026e2e1..85fff75 100644
--- a/src/SplitClient.tsx
+++ b/src/SplitClient.tsx
@@ -25,6 +25,7 @@ export class SplitComponent extends React.Component<IUpdateProps & { factory: Sp
   // But it implies to have another instance property to use instead of the state, because we need a unique reference value for SplitContext.Producer
   static getDerivedStateFromProps(props: ISplitClientProps & { factory: SplitIO.IBrowserSDK | null, client: SplitIO.IBrowserClient | null }, state: ISplitContextValues) {
     const { client, factory, attributes } = props;
+    // @TODO check if apply over client or over stage.client
     if (attributes)
       client?.setAttributes(attributes);
     else
diff --git a/src/SplitTreatments.tsx b/src/SplitTreatments.tsx
index 7bd118c..d7b6def 100644
--- a/src/SplitTreatments.tsx
+++ b/src/SplitTreatments.tsx
@@ -9,10 +9,11 @@ function argsAreEqual(newArgs: any[], lastArgs: any[]): boolean {
   return newArgs[0] === lastArgs[0] && // client
     newArgs[1] === lastArgs[1] && // lastUpdate
     shallowEqual(newArgs[2], lastArgs[2]) && // names
-    shallowEqual(newArgs[3], lastArgs[3]); // attributes
+    shallowEqual(newArgs[3], lastArgs[3]) && // attributes
+    shallowEqual(newArgs[4], lastArgs[4]); // client attributes
 }
 
-function evaluateSplits(client: SplitIO.IBrowserClient, lastUpdate: number, names: SplitIO.SplitNames, attributes?: SplitIO.Attributes) {
+function evaluateSplits(client: SplitIO.IBrowserClient, lastUpdate: number, names: SplitIO.SplitNames, attributes?: SplitIO.Attributes, clientAttributes?: SplitIO.Attributes) {
   return client.getTreatmentsWithConfig(names, attributes);
 }
 
@@ -40,7 +41,7 @@ class SplitTreatments extends React.Component<ISplitTreatmentsProps> {
           let treatments;
           const isOperational = !isDestroyed && (isReady || isReadyFromCache);
           if (client && isOperational) {
-            treatments = this.evaluateSplits(client, lastUpdate, names, attributes);
+            treatments = this.evaluateSplits(client, lastUpdate, names, attributes, client.getAttributes());
           } else {
             treatments = getControlTreatmentsWithConfig(names);
             if (!client) { this.logWarning = true; }
diff --git a/src/__tests__/SplitClient.test.tsx b/src/__tests__/SplitClient.test.tsx
index 2ffb57d..9f0cb2d 100644
--- a/src/__tests__/SplitClient.test.tsx
+++ b/src/__tests__/SplitClient.test.tsx
@@ -1,5 +1,7 @@
 import React from 'react';
 import { mount, ReactWrapper } from 'enzyme';
+// @ts-ignore
+import { SplitFactory as originalSplitFactory } from '../../lib/splitio/index';
 
 /** Mocks and test utils */
 import { mockSdk, Event, assertNoListeners, clientListenerCount } from './testUtils/mockSplitSdk';
@@ -15,6 +17,7 @@ import SplitFactory from '../SplitFactory';
 import SplitClient from '../SplitClient';
 import SplitContext, { ISplitContextValues } from '../SplitContext';
 import { ERROR_SC_NO_FACTORY } from '../constants';
+import SplitIO from '@splitsoftware/splitio/types/splitio';
 
 describe('SplitClient', () => {
 
@@ -392,4 +395,77 @@ describe('SplitClient', () => {
       </SplitFactory>);
   });
 
+  test('calls client setAttributes and clearAttributes.', (done) => {
+    let renderTimes = 0;
+    const attributesFactory = { at1: 'at1' };
+    const attributesClient = { at2: 'at2' };
+    const factory = originalSplitFactory({
+      core: {
+        authorizationKey: 'localhost',
+        key: 'emma'
+      },
+      features: {
+        test_split: 'on'
+      }
+    })
+
+    const mainClient = factory?.client();
+    const mainClientSpy = {
+      setAttributes: jest.spyOn(mainClient, 'setAttributes'),
+      clearAttributes: jest.spyOn(mainClient, 'clearAttributes')
+    }
+
+    const client = factory?.client('user1');
+    const clientSpy = {
+      setAttributes: jest.spyOn(client, 'setAttributes'),
+      clearAttributes: jest.spyOn(client, 'clearAttributes'),
+    }
+
+    function Component({ attributesFactory, attributesClient }: { attributesFactory: SplitIO.Attributes, attributesClient: SplitIO.Attributes }) {
+      return (
+        <SplitFactory factory={factory} attributes={attributesFactory} >
+          <SplitClient splitKey='user1' attributes={attributesClient} >
+            {({ client }: ISplitClientChildProps) => {
+              renderTimes++;
+              switch (renderTimes) {
+                case 1:
+                  expect(mainClientSpy.setAttributes).lastCalledWith(attributesFactory);
+                  expect(clientSpy.setAttributes).lastCalledWith(attributesClient);
+                  expect(mainClient?.getAttributes()).toStrictEqual(attributesFactory);
+                  expect(client?.getAttributes()).toStrictEqual(attributesClient);
+                  break;
+                case 2:
+                  expect(mainClientSpy.clearAttributes).toBeCalledTimes(1);
+                  expect(clientSpy.setAttributes).lastCalledWith(attributesClient);
+                  expect(mainClient?.getAttributes()).toStrictEqual({});
+                  expect(client?.getAttributes()).toStrictEqual({at2: 'at2', at3:'at3'});
+                  break;
+                case 3:
+                  expect(mainClientSpy.setAttributes).lastCalledWith({at4: 'at4'});
+                  expect(clientSpy.clearAttributes).toBeCalledTimes(2);
+                  expect(mainClient?.getAttributes()).toStrictEqual({at4: 'at4'});
+                  expect(client?.getAttributes()).toStrictEqual({});
+                  break;
+                case 4:
+                  expect(mainClientSpy.clearAttributes).toBeCalledTimes(2);
+                  expect(clientSpy.clearAttributes).toBeCalledTimes(4);
+                  expect(mainClient?.getAttributes()).toStrictEqual({});
+                  expect(client?.getAttributes()).toStrictEqual({});
+                  done();
+                }
+              return null;
+            }}
+          </SplitClient>
+        </SplitFactory>
+      );
+    }
+
+    const wrapper = mount(<Component attributesFactory={attributesFactory} attributesClient={attributesClient} />);
+
+    wrapper.setProps({ attributesFactory: undefined, attributesClient: {at3: 'at3'} });
+    wrapper.setProps({ attributesFactory: {at4: 'at4'}, attributesClient: undefined });
+    wrapper.setProps({ attributesFactory: undefined, attributesClient: undefined });
+
+  });
+
 });
diff --git a/src/__tests__/testUtils/mockSplitSdk.ts b/src/__tests__/testUtils/mockSplitSdk.ts
index 161630f..56b82a7 100644
--- a/src/__tests__/testUtils/mockSplitSdk.ts
+++ b/src/__tests__/testUtils/mockSplitSdk.ts
@@ -54,6 +54,15 @@ function mockClient(key: SplitIO.SplitKey, trafficType?: string) {
   const getTreatmentsWithConfig: jest.Mock = jest.fn(() => {
     return 'getTreatmentsWithConfig';
   });
+  const setAttributes: jest.Mock = jest.fn(() => {
+    return true;
+  });
+  const clearAttributes: jest.Mock = jest.fn(() => {
+    return true;
+  });
+  const getAttributes: jest.Mock = jest.fn(() => {
+    return true;
+  });
   const ready: jest.Mock = jest.fn(() => {
     return new Promise((res, rej) => {
       if (__isReady__) res();
@@ -94,6 +103,9 @@ function mockClient(key: SplitIO.SplitKey, trafficType?: string) {
     ready,
     destroy,
     Event,
+    setAttributes,
+    clearAttributes,
+    getAttributes,
     // EventEmitter exposed to trigger events manually
     __emitter__,
     // Count of internal listeners set by the client mock, used to assert how many external listeners were attached
diff --git a/src/__tests__/useClient.test.tsx b/src/__tests__/useClient.test.tsx
index d677546..daf66ab 100644
--- a/src/__tests__/useClient.test.tsx
+++ b/src/__tests__/useClient.test.tsx
@@ -1,6 +1,9 @@
 import React from 'react';
 import { mount, shallow } from 'enzyme';
 
+// @ts-ignore
+import { SplitFactory as originalSplitFactory } from '../../lib/splitio/index';
+
 /** Mocks */
 import { mockSdk } from './testUtils/mockSplitSdk';
 jest.mock('@splitsoftware/splitio', () => {
@@ -76,4 +79,76 @@ describe('useClient', () => {
     expect(sharedClient).toBe(null);
   });
 
+  test('calls client setAttributes and clearAttributes.', (done) => {
+
+    let renderTimes = 0;
+    const attributesFactory = { at1: 'at1' };
+    const attributesClient = { at2: 'at2' };
+
+    const factory = originalSplitFactory({
+      core: {
+        authorizationKey: 'localhost',
+        key: 'emma'
+      },
+      features: {
+        test_split: 'on'
+      }
+    })
+
+    const mainClient = factory.client();
+    const mainClientSpy = {
+      setAttributes: jest.spyOn(mainClient, 'setAttributes'),
+      clearAttributes: jest.spyOn(mainClient, 'clearAttributes')
+    }
+
+    const client = factory.client('user1', 'user');
+    const clientSpy = {
+      setAttributes: jest.spyOn(client, 'setAttributes'),
+      clearAttributes: jest.spyOn(client, 'clearAttributes'),
+    }
+
+    function Component({ attributesFactory, attributesClient }: { attributesFactory: any, attributesClient: any }) {
+      return (
+        <SplitFactory factory={factory} attributes={attributesFactory} >{
+          React.createElement(() => {
+            useClient('user1', 'user', attributesClient);
+            renderTimes++;
+            switch (renderTimes) {
+              case 1:
+                expect(mainClientSpy.setAttributes).lastCalledWith(attributesFactory);
+                expect(clientSpy.setAttributes).lastCalledWith(attributesClient);
+                expect(mainClient?.getAttributes()).toStrictEqual(attributesFactory);
+                expect(client?.getAttributes()).toStrictEqual(attributesClient);
+                break;
+              case 2:
+                expect(mainClientSpy.clearAttributes).toBeCalledTimes(1);
+                expect(clientSpy.setAttributes).lastCalledWith(attributesClient);
+                expect(mainClient?.getAttributes()).toStrictEqual({});
+                expect(client?.getAttributes()).toStrictEqual({at2: 'at2', at3:'at3'});
+                break;
+              case 3:
+                expect(mainClientSpy.setAttributes).lastCalledWith({at4: 'at4'});
+                expect(clientSpy.clearAttributes).toBeCalledTimes(1);
+                expect(mainClient?.getAttributes()).toStrictEqual({at4: 'at4'});
+                expect(client?.getAttributes()).toStrictEqual({});
+                break;
+              case 4:
+                expect(mainClientSpy.clearAttributes).toBeCalledTimes(2);
+                expect(clientSpy.clearAttributes).toBeCalledTimes(2);
+                expect(mainClient?.getAttributes()).toStrictEqual({});
+                expect(client?.getAttributes()).toStrictEqual({});
+                done()
+              }
+            return null
+        })}</SplitFactory>
+      )
+    }
+
+    const wrapper = mount(<Component attributesFactory={attributesFactory} attributesClient={attributesClient} />);
+
+    wrapper.setProps({ attributesFactory: undefined, attributesClient: {at3: 'at3'} });
+    wrapper.setProps({ attributesFactory: {at4: 'at4'}, attributesClient: undefined });
+    wrapper.setProps({ attributesFactory: undefined, attributesClient: undefined });
+  });
+
 });
diff --git a/src/__tests__/withSplitClient.test.tsx b/src/__tests__/withSplitClient.test.tsx
index 2ce10d0..124b9c0 100644
--- a/src/__tests__/withSplitClient.test.tsx
+++ b/src/__tests__/withSplitClient.test.tsx
@@ -1,5 +1,7 @@
 import React from 'react';
 import { mount, shallow } from 'enzyme';
+// @ts-ignore
+import { SplitFactory as originalSplitFactory } from '../../lib/splitio/index';
 
 /** Mocks */
 import { mockSdk, Event } from './testUtils/mockSplitSdk';
@@ -72,4 +74,76 @@ describe('SplitClient', () => {
     expect(wrapper.prop('updateOnSdkReadyFromCache')).toBe(updateOnSdkReadyFromCache);
   });
 
+  test('passes attributes to Splitclient', () => {
+
+    let renderTimes = 0;
+
+    const attributesFactory = { at1: 'at1' };
+    const attributesClient = { at2: 'at2' };
+
+    const factory = originalSplitFactory({
+      core: {
+        authorizationKey: 'localhost',
+        key: 'emma'
+      },
+      features: {
+        test_split: 'on'
+      }
+    })
+
+    const mainClient = factory?.client();
+    const mainClientSpy = {
+      setAttributes: jest.spyOn(mainClient, 'setAttributes'),
+      clearAttributes: jest.spyOn(mainClient, 'clearAttributes')
+    }
+
+    const client = factory?.client('user1', 'user');
+    const clientSpy = {
+      setAttributes: jest.spyOn(client, 'setAttributes'),
+      clearAttributes: jest.spyOn(client, 'clearAttributes'),
+    }
+
+    function Component({ attributesFactory, attributesClient }: { attributesFactory: any, attributesClient: any }) {
+      const Component = withSplitFactory(undefined, factory, attributesFactory)(
+      withSplitClient('user1', 'user', attributesClient)(
+        () => {
+          renderTimes++;
+          switch (renderTimes) {
+            case 1:
+              expect(mainClientSpy.setAttributes).lastCalledWith(attributesFactory);
+              expect(clientSpy.setAttributes).lastCalledWith(attributesClient);
+              expect(mainClient?.getAttributes()).toStrictEqual(attributesFactory);
+              expect(client?.getAttributes()).toStrictEqual(attributesClient);
+              break;
+            case 2:
+              expect(mainClientSpy.clearAttributes).toBeCalledTimes(1);
+              expect(clientSpy.setAttributes).lastCalledWith(attributesClient);
+              expect(mainClient?.getAttributes()).toStrictEqual({});
+              expect(client?.getAttributes()).toStrictEqual({at2: 'at2', at3:'at3'});
+              break;
+            case 3:
+              expect(mainClientSpy.setAttributes).lastCalledWith({at4: 'at4'});
+              expect(clientSpy.clearAttributes).toBeCalledTimes(1);
+              expect(mainClient?.getAttributes()).toStrictEqual({at4: 'at4'});
+              expect(client?.getAttributes()).toStrictEqual({});
+              break;
+            case 4:
+              expect(mainClientSpy.clearAttributes).toBeCalledTimes(2);
+              expect(clientSpy.clearAttributes).toBeCalledTimes(2);
+              expect(mainClient?.getAttributes()).toStrictEqual({});
+              expect(client?.getAttributes()).toStrictEqual({});
+
+            }
+          return null;
+        }))
+        return Component}
+
+    // @ts-ignore
+    const wrapper = mount(<Component />);
+
+    wrapper.setProps({ attributesFactory: undefined, attributesClient: {at3: 'at3'} });
+    wrapper.setProps({ attributesFactory: {at4: 'at4'}, attributesClient: undefined });
+    wrapper.setProps({ attributesFactory: undefined, attributesClient: undefined });
+  });
+
 });
diff --git a/src/useClient.ts b/src/useClient.ts
index 26fb230..2383395 100644
--- a/src/useClient.ts
+++ b/src/useClient.ts
@@ -14,13 +14,13 @@ import { getSplitSharedClient, checkHooks } from './utils';
 const useClient = (key?: SplitIO.SplitKey, trafficType?: string, attributes?: SplitIO.Attributes): SplitIO.IBrowserClient | null => {
   if (!checkHooks(ERROR_UC_NO_USECONTEXT)) return null;
   const { factory, client } = React.useContext(SplitContext);
+  if (key) {
+    return factory ? getSplitSharedClient(factory, key, trafficType, attributes) : null;
+  }
   if (attributes)
     client?.setAttributes(attributes);
   else
     client?.clearAttributes();
-  if (key) {
-    return factory ? getSplitSharedClient(factory, key, trafficType, attributes) : null;
-  }
   return client;
 };