Support custom data processing during import

Summary: Per title, this allows for pre-processing data after it is deserialized and before it is stored in the plugin

Reviewed By: nikoant

Differential Revision: D26126423

fbshipit-source-id: bc08a6ab205d2a0d551515563cd85a197595ddb2

desktop/app/src/utils/__tests__/exportData.node.tsx

@@ -1380,3 +1380,141 @@ test('Sandy device plugins with custom export are export properly', async () =>
     [sandyDeviceTestPlugin.id]: {customExport: true},
   });
 });
+
+test('Sandy plugin with custom import', async () => {
+  const plugin = new _SandyPluginDefinition(
+    TestUtils.createMockPluginDetails(),
+    {
+      plugin(client: PluginClient) {
+        const counter = createState(0);
+        client.onImport((data) => {
+          counter.set(data.count);
+        });
+
+        return {
+          counter,
+        };
+      },
+      Component() {
+        return null;
+      },
+    },
+  );
+
+  const {store} = await renderMockFlipperWithPlugin(plugin);
+
+  const data = {
+    clients: [
+      {
+        id:
+          'TestApp#Android#MockAndroidDevice#2e52cea6-94b0-4ea1-b9a8-c9135ede14ca-serial',
+        query: {
+          app: 'TestApp',
+          device: 'MockAndroidDevice',
+          device_id: '2e52cea6-94b0-4ea1-b9a8-c9135ede14ca-serial',
+          os: 'Android',
+          sdk_version: 4,
+        },
+      },
+    ],
+    device: {
+      deviceType: 'physical',
+      logs: [],
+      os: 'Android',
+      serial: '2e52cea6-94b0-4ea1-b9a8-c9135ede14ca-serial',
+      title: 'MockAndroidDevice',
+    },
+    deviceScreenshot: null,
+    fileVersion: '0.9.99',
+    flipperReleaseRevision: undefined,
+    pluginStates2: {
+      'TestApp#Android#MockAndroidDevice#2e52cea6-94b0-4ea1-b9a8-c9135ede14ca-serial': {
+        [plugin.id]: {
+          count: 4,
+        },
+      },
+    },
+    store: {
+      activeNotifications: [],
+      pluginStates: {},
+    },
+  };
+
+  await importDataToStore('unittest.json', JSON.stringify(data), store);
+
+  expect(
+    store
+      .getState()
+      .connections.clients[0].sandyPluginStates.get(plugin.id)
+      ?.instanceApi.counter.get(),
+  ).toBe(0);
+  expect(
+    store
+      .getState()
+      .connections.clients[1].sandyPluginStates.get(plugin.id)
+      ?.instanceApi.counter.get(),
+  ).toBe(4);
+});
+
+test('Sandy device plugin with custom import', async () => {
+  const plugin = new _SandyPluginDefinition(
+    TestUtils.createMockPluginDetails(),
+    {
+      supportsDevice: () => true,
+      devicePlugin(client: DevicePluginClient) {
+        const counter = createState(0);
+        client.onImport((data) => {
+          counter.set(data.count);
+        });
+
+        return {
+          counter,
+        };
+      },
+      Component() {
+        return null;
+      },
+    },
+  );
+
+  const data = {
+    clients: [],
+    device: {
+      deviceType: 'archivedPhysical',
+      logs: [],
+      os: 'Android',
+      serial: '2e52cea6-94b0-4ea1-b9a8-c9135ede14ca-serial',
+      title: 'MockAndroidDevice',
+      pluginStates: {
+        [plugin.id]: {
+          count: 2,
+        },
+      },
+    },
+    deviceScreenshot: null,
+    fileVersion: '0.9.99',
+    flipperReleaseRevision: undefined,
+    pluginStates2: {},
+    store: {
+      activeNotifications: [],
+      pluginStates: {},
+    },
+  };
+
+  const {store} = await renderMockFlipperWithPlugin(plugin);
+
+  await importDataToStore('unittest.json', JSON.stringify(data), store);
+
+  expect(
+    store
+      .getState()
+      .connections.devices[0].sandyPluginStates.get(plugin.id)
+      ?.instanceApi.counter.get(),
+  ).toBe(0);
+  expect(
+    store
+      .getState()
+      .connections.devices[1].sandyPluginStates.get(plugin.id)
+      ?.instanceApi.counter.get(),
+  ).toBe(2);
+});

desktop/app/src/utils/exportData.tsx

@@ -329,7 +329,6 @@ async function addSaltToDeviceSerial({
   selectedPlugins,
   pluginStates2,
   devicePluginStates,
-  idler,
 }: AddSaltToDeviceSerialOptions): Promise<ExportType> {
   const {serial} = device;
   const newSerial = salt + '-' + serial;

desktop/flipper-plugin/package.json

@@ -16,6 +16,7 @@
   },
   "devDependencies": {
     "@types/jest": "^26.0.3",
+    "jest-mock-console": "^1.0.1",
     "typescript": "^4.1.2"
   },
   "peerDependencies": {

desktop/flipper-plugin/src/__tests__/test-utils-device.node.tsx

@@ -128,16 +128,13 @@ test('device plugins support non-serializable state', async () => {
 });
 
 test('device plugins support restoring state', async () => {
-  const {exportState} = TestUtils.startPlugin(
+  const {exportState, instance} = TestUtils.startPlugin(
     {
       plugin() {
         const field1 = createState(1, {persist: 'field1'});
         const field2 = createState(2);
         const field3 = createState(3, {persist: 'field3'});
-        expect(field1.get()).toBe('a');
+        return {field1, field2, field3};
-        expect(field2.get()).toBe(2);
-        expect(field3.get()).toBe('b');
-        return {};
       },
       Component() {
         return null;
@@ -147,5 +144,10 @@ test('device plugins support restoring state', async () => {
       initialState: {field1: 'a', field3: 'b'},
     },
   );
+
+  const {field1, field2, field3} = instance;
+  expect(field1.get()).toBe('a');
+  expect(field2.get()).toBe(2);
+  expect(field3.get()).toBe('b');
   expect(exportState()).toEqual({field1: 'a', field3: 'b'});
 });

desktop/flipper-plugin/src/__tests__/test-utils.node.tsx

@@ -12,6 +12,8 @@ import * as testPlugin from './TestPlugin';
 import {createState} from '../state/atom';
 import {PluginClient} from '../plugin/Plugin';
 import {DevicePluginClient} from '../plugin/DevicePlugin';
+import mockConsole from 'jest-mock-console';
+import {sleep} from '../utils/sleep';
 
 test('it can start a plugin and lifecycle events', () => {
   const {instance, ...p} = TestUtils.startPlugin(testPlugin);
@@ -217,16 +219,17 @@ test('plugins support non-serializable state', async () => {
 });
 
 test('plugins support restoring state', async () => {
-  const {exportState} = TestUtils.startPlugin(
+  const {exportState, instance} = TestUtils.startPlugin(
     {
       plugin() {
         const field1 = createState(1, {persist: 'field1'});
         const field2 = createState(2);
         const field3 = createState(3, {persist: 'field3'});
-        expect(field1.get()).toBe('a');
+        return {
-        expect(field2.get()).toBe(2);
+          field1,
-        expect(field3.get()).toBe('b');
+          field2,
-        return {};
+          field3,
+        };
       },
       Component() {
         return null;
@@ -236,6 +239,12 @@ test('plugins support restoring state', async () => {
       initialState: {field1: 'a', field3: 'b'},
     },
   );
+
+  const {field1, field2, field3} = instance;
+  expect(field1.get()).toBe('a');
+  expect(field2.get()).toBe(2);
+  expect(field3.get()).toBe('b');
+
   expect(exportState()).toEqual({field1: 'a', field3: 'b'});
 });
 
@@ -256,6 +265,125 @@ test('plugins cannot use a persist key twice', async () => {
   );
 });
 
+test('plugins can have custom import handler', () => {
+  const {instance} = TestUtils.startPlugin(
+    {
+      plugin(client: PluginClient) {
+        const field1 = createState(0);
+        const field2 = createState(0);
+
+        client.onImport((data) => {
+          field1.set(data.a);
+          field2.set(data.b);
+        });
+
+        return {field1, field2};
+      },
+      Component() {
+        return null;
+      },
+    },
+    {
+      initialState: {
+        a: 1,
+        b: 2,
+      },
+    },
+  );
+  expect(instance.field1.get()).toBe(1);
+  expect(instance.field2.get()).toBe(2);
+});
+
+test('plugins cannot combine import handler with persist option', async () => {
+  expect(() => {
+    TestUtils.startPlugin({
+      plugin(client: PluginClient) {
+        const field1 = createState(1, {persist: 'f1'});
+        const field2 = createState(1, {persist: 'f2'});
+        client.onImport(() => {});
+        return {field1, field2};
+      },
+      Component() {
+        return null;
+      },
+    });
+  }).toThrowErrorMatchingInlineSnapshot(
+    `"A custom onImport handler was defined for plugin 'TestPlugin', the 'persist' option of states f1, f2 should not be set."`,
+  );
+});
+
+test('plugins can handle import errors', async () => {
+  const restoreConsole = mockConsole();
+  let instance: any;
+  try {
+    instance = TestUtils.startPlugin(
+      {
+        plugin(client: PluginClient) {
+          const field1 = createState(0);
+          const field2 = createState(0);
+
+          client.onImport(() => {
+            throw new Error('Oops');
+          });
+
+          return {field1, field2};
+        },
+        Component() {
+          return null;
+        },
+      },
+      {
+        initialState: {
+          a: 1,
+          b: 2,
+        },
+      },
+    ).instance;
+    // @ts-ignore
+    expect(console.error.mock.calls).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          "Error occurred when importing date for plugin 'TestPlugin': 'Error: Oops",
+          [Error: Oops],
+        ],
+      ]
+    `);
+  } finally {
+    restoreConsole();
+  }
+  expect(instance.field1.get()).toBe(0);
+  expect(instance.field2.get()).toBe(0);
+});
+
+test('plugins can have custom export handler', async () => {
+  const {exportStateAsync} = TestUtils.startPlugin(
+    {
+      plugin(client: PluginClient) {
+        const field1 = createState(0, {persist: 'field1'});
+
+        client.onExport(async () => {
+          await sleep(10);
+          return {
+            b: 3,
+          };
+        });
+
+        return {field1};
+      },
+      Component() {
+        return null;
+      },
+    },
+    {
+      initialState: {
+        a: 1,
+        b: 2,
+      },
+    },
+  );
+  expect(await exportStateAsync()).toEqual({b: 3});
+});
+
 test('plugins can receive deeplinks', async () => {
   const plugin = TestUtils.startPlugin({
     plugin(client: PluginClient) {

desktop/flipper-plugin/src/plugin/PluginBase.tsx

@@ -15,11 +15,13 @@ import {FlipperLib} from './FlipperLib';
 import {Device, RealFlipperDevice} from './DevicePlugin';
 import {batched} from '../state/batch';
 import {Idler} from '../utils/Idler';
+import {message} from 'antd';
 
 type StateExportHandler = (
   idler: Idler,
   onStatusMessage: (msg: string) => void,
 ) => Promise<Record<string, any>>;
+type StateImportHandler = (data: Record<string, any>) => void;
 
 export interface BasePluginClient {
   readonly device: Device;
@@ -50,6 +52,12 @@ export interface BasePluginClient {
    */
   onExport(exporter: StateExportHandler): void;
 
+  /**
+   * Triggered directly after the plugin instance was created, if the plugin is being restored from a snapshot.
+   * Should be the inverse of the onExport handler
+   */
+  onImport(handler: StateImportHandler): void;
+
   /**
    * Register menu entries in the Flipper toolbar
    */
@@ -96,12 +104,15 @@ export abstract class BasePluginInstance {
 
   // temporarily field that is used during deserialization
   initialStates?: Record<string, any>;
+
   // all the atoms that should be serialized when making an export / import
   rootStates: Record<string, Atom<any>> = {};
   // last seen deeplink
   lastDeeplink?: any;
   // export handler
   exportHandler?: StateExportHandler;
+  // import handler
+  importHandler?: StateImportHandler;
 
   menuEntries: NormalizedMenuEntry[] = [];
 
@@ -139,6 +150,37 @@ export abstract class BasePluginInstance {
     try {
       this.instanceApi = batched(factory)();
     } finally {
+      // check if we have both an import handler and rootStates; probably dev error
+      if (this.importHandler && Object.keys(this.rootStates).length > 0) {
+        throw new Error(
+          `A custom onImport handler was defined for plugin '${
+            this.definition.id
+          }', the 'persist' option of states ${Object.keys(
+            this.rootStates,
+          ).join(', ')} should not be set.`,
+        );
+      }
+      if (this.initialStates) {
+        if (this.importHandler) {
+          try {
+            this.importHandler(this.initialStates);
+          } catch (e) {
+            const msg = `Error occurred when importing date for plugin '${this.definition.id}': '${e}`;
+            console.error(msg, e);
+            message.error(msg);
+          }
+        } else {
+          for (const key in this.rootStates) {
+            if (key in this.initialStates) {
+              this.rootStates[key].set(this.initialStates[key]);
+            } else {
+              console.warn(
+                `Tried to initialize plugin with existing data, however data for "${key}" is missing. Was the export created with a different Flipper version?`,
+              );
+            }
+          }
+        }
+      }
       this.initialStates = undefined;
       setCurrentPluginInstance(undefined);
     }
@@ -165,6 +207,12 @@ export abstract class BasePluginInstance {
         }
         this.exportHandler = cb;
       },
+      onImport: (cb) => {
+        if (this.importHandler) {
+          throw new Error('onImport handler already set');
+        }
+        this.importHandler = cb;
+      },
       addMenuEntry: (...entries) => {
         for (const entry of entries) {
           const normalized = normalizeMenuEntry(entry);

desktop/flipper-plugin/src/state/atom.tsx

@@ -73,16 +73,7 @@ export function createState<T>(
 ): Atom<T> {
   const atom = new AtomValue<T>(initialValue);
   if (getCurrentPluginInstance() && options.persist) {
-    const {initialStates, rootStates} = getCurrentPluginInstance()!;
+    const {rootStates} = getCurrentPluginInstance()!;
-    if (initialStates) {
-      if (options.persist in initialStates) {
-        atom.set(initialStates[options.persist]);
-      } else {
-        console.warn(
-          `Tried to initialize plugin with existing data, however data for "${options.persist}" is missing. Was the export created with a different Flipper version?`,
-        );
-      }
-    }
     if (rootStates[options.persist]) {
       throw new Error(
         `Some other state is already persisting with key "${options.persist}"`,

desktop/yarn.lock

@@ -7736,6 +7736,11 @@ jest-message-util@^26.6.0:
     slash "^3.0.0"
     stack-utils "^2.0.2"
 
+jest-mock-console@^1.0.1:
+  version "1.0.1"
+  resolved "https://registry.yarnpkg.com/jest-mock-console/-/jest-mock-console-1.0.1.tgz#07978047735a782d0d4172d1afcabd82f6de9b08"
+  integrity sha512-Bn+Of/cvz9LOEEeEg5IX5Lsf8D2BscXa3Zl5+vSVJl37yiT8gMAPPKfE09jJOwwu1zbagL11QTrH+L/Gn8udOg==
+
 jest-mock@^25.0.0, jest-mock@^25.1.0, jest-mock@^25.5.0:
   version "25.5.0"
   resolved "https://registry.yarnpkg.com/jest-mock/-/jest-mock-25.5.0.tgz#a91a54dabd14e37ecd61665d6b6e06360a55387a"

docs/extending/flipper-plugin.mdx

@@ -142,6 +142,13 @@ Usage: `client.onExport(callback: (idler, onStatusMessage) => Promise<state>)`
 Overrides the default serialization behavior of this plugin. Should return a promise with state that is to be stored. 
 This process is async, so it is possible to first fetch some additional state from the device. 
 
+#### `onImport`
+
+Usage: `client.onImport(callback: (snapshot) => void)`
+
+Overrides the default de-serialization behavior of this plugin. Use it to update the state based on the snapshot data.
+This hook will be called immediately after constructing the plugin instance.
+
 ### Methods
 
 #### `send`