2022 ExperimentProvider.tsx

ExperimentProvider.tsx

// ExperimentProvider.tsx
import { FC, useMemo, useState } from 'react';
import {
  AmplitudeUserProvider,
  Experiment,
  ExperimentClient,
} from '@amplitude/experiment-js-client';
import Amplitude from 'amplitude-js';

import {
  AmplitudeExperimentContext,
  AmplitudeExperimentContextValue,
} from './ExperimentContext';

export interface AmplitudeExperimentProviderProps {
  envExperimentKey: string;
  amplitudeApiKey: string;
  user?: { originalId: number };
  debugMode?: boolean;
}

const AmplitudeExperimentProvider: FC<AmplitudeExperimentProviderProps> = ({
  envExperimentKey,
  amplitudeApiKey,
  debugMode = false,
  user,
  children,
}) => {
  const [, setLoaded] = useState(false);

  const startExperiment = async (experiment: ExperimentClient) => {
    await experiment.fetch();
    setLoaded(true);
  };

  const experiment = useMemo<ExperimentClient | undefined>(() => {
    const amplitude = Amplitude.getInstance();

    if (user?.originalId) {
      amplitude.setUserId(user.originalId.toString());
    }

    amplitude.init(amplitudeApiKey);

    const internExperiment = Experiment.initialize(envExperimentKey, {
      userProvider: new AmplitudeUserProvider(amplitude),
      debug: debugMode,
    });

    startExperiment(internExperiment);

    console.log('Experiment initialized');

    return internExperiment;
  }, [debugMode, amplitudeApiKey, envExperimentKey, user]);

  const isVariantActive = (variantId: string, variantValue?: string) => {
    const variantObject = experiment.variant(variantId);

    if (variantValue) {
      return variantObject.value === variantValue;
    }

    return variantObject.value !== 'control';
  };

  const isControlActive = (
    variantId: string,
    defaultControl: string = 'control',
  ) => isVariantActive(variantId, defaultControl);

  function getPayloadVariant<T>(variantId: string): T | null {
    const variantObject = experiment.variant(variantId);

    return variantObject.payload as T;
  }

  const value: AmplitudeExperimentContextValue = {
    experiment,
    isVariantActive,
    isControlActive,
    getPayloadVariant,
  };

  // Avoid children re-render if state provider is changed.
  return useMemo(
    () => (
      <AmplitudeExperimentContext.Provider value={value}>
        {children}
      </AmplitudeExperimentContext.Provider>
    ),
    [children],
  );
};

export { AmplitudeExperimentContext, AmplitudeExperimentProvider };