Skip to content

Commit 591f7f0

Browse files
committed
Add BaseControllerV2 schema
A schema has been added to the new BaseController class as a required constructor parameter. The schema describes which pieces of state should be persisted, and how to get an 'anonymized' snapshot of the controller state. This is part of the controller redesign (#337).
1 parent 6e0b2da commit 591f7f0

File tree

2 files changed

+170
-13
lines changed

2 files changed

+170
-13
lines changed

src/BaseControllerV2.test.ts

+127-12
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,19 @@
11
import type { Draft } from 'immer';
22
import * as sinon from 'sinon';
33

4-
import { BaseController } from './BaseControllerV2';
4+
import { BaseController, getAnonymizedState, getPersistentState } from './BaseControllerV2';
55

66
interface MockControllerState {
77
count: number;
88
}
99

10+
const mockControllerSchema = {
11+
count: {
12+
persist: true,
13+
anonymous: true,
14+
},
15+
};
16+
1017
class MockController extends BaseController<MockControllerState> {
1118
update(callback: (state: Draft<MockControllerState>) => void | MockControllerState) {
1219
super.update(callback);
@@ -19,21 +26,27 @@ class MockController extends BaseController<MockControllerState> {
1926

2027
describe('BaseController', () => {
2128
it('should set initial state', () => {
22-
const controller = new MockController({ count: 0 });
29+
const controller = new MockController({ count: 0 }, mockControllerSchema);
2330

2431
expect(controller.state).toEqual({ count: 0 });
2532
});
2633

34+
it('should set initial schema', () => {
35+
const controller = new MockController({ count: 0 }, mockControllerSchema);
36+
37+
expect(controller.schema).toEqual(mockControllerSchema);
38+
});
39+
2740
it('should not allow mutating state directly', () => {
28-
const controller = new MockController({ count: 0 });
41+
const controller = new MockController({ count: 0 }, mockControllerSchema);
2942

3043
expect(() => {
3144
controller.state = { count: 1 };
3245
}).toThrow();
3346
});
3447

3548
it('should allow updating state by modifying draft', () => {
36-
const controller = new MockController({ count: 0 });
49+
const controller = new MockController({ count: 0 }, mockControllerSchema);
3750

3851
controller.update((draft) => {
3952
draft.count += 1;
@@ -43,7 +56,7 @@ describe('BaseController', () => {
4356
});
4457

4558
it('should allow updating state by return a value', () => {
46-
const controller = new MockController({ count: 0 });
59+
const controller = new MockController({ count: 0 }, mockControllerSchema);
4760

4861
controller.update(() => {
4962
return { count: 1 };
@@ -53,7 +66,7 @@ describe('BaseController', () => {
5366
});
5467

5568
it('should throw an error if update callback modifies draft and returns value', () => {
56-
const controller = new MockController({ count: 0 });
69+
const controller = new MockController({ count: 0 }, mockControllerSchema);
5770

5871
expect(() => {
5972
controller.update((draft) => {
@@ -64,7 +77,7 @@ describe('BaseController', () => {
6477
});
6578

6679
it('should inform subscribers of state changes', () => {
67-
const controller = new MockController({ count: 0 });
80+
const controller = new MockController({ count: 0 }, mockControllerSchema);
6881
const listener1 = sinon.stub();
6982
const listener2 = sinon.stub();
7083

@@ -81,7 +94,7 @@ describe('BaseController', () => {
8194
});
8295

8396
it('should inform a subscriber of each state change once even after multiple subscriptions', () => {
84-
const controller = new MockController({ count: 0 });
97+
const controller = new MockController({ count: 0 }, mockControllerSchema);
8598
const listener1 = sinon.stub();
8699

87100
controller.subscribe(listener1);
@@ -95,7 +108,7 @@ describe('BaseController', () => {
95108
});
96109

97110
it('should no longer inform a subscriber about state changes after unsubscribing', () => {
98-
const controller = new MockController({ count: 0 });
111+
const controller = new MockController({ count: 0 }, mockControllerSchema);
99112
const listener1 = sinon.stub();
100113

101114
controller.subscribe(listener1);
@@ -108,7 +121,7 @@ describe('BaseController', () => {
108121
});
109122

110123
it('should no longer inform a subscriber about state changes after unsubscribing once, even if they subscribed many times', () => {
111-
const controller = new MockController({ count: 0 });
124+
const controller = new MockController({ count: 0 }, mockControllerSchema);
112125
const listener1 = sinon.stub();
113126

114127
controller.subscribe(listener1);
@@ -122,7 +135,7 @@ describe('BaseController', () => {
122135
});
123136

124137
it('should allow unsubscribing listeners who were never subscribed', () => {
125-
const controller = new MockController({ count: 0 });
138+
const controller = new MockController({ count: 0 }, mockControllerSchema);
126139
const listener1 = sinon.stub();
127140

128141
expect(() => {
@@ -131,7 +144,7 @@ describe('BaseController', () => {
131144
});
132145

133146
it('should no longer update subscribers after being destroyed', () => {
134-
const controller = new MockController({ count: 0 });
147+
const controller = new MockController({ count: 0 }, mockControllerSchema);
135148
const listener1 = sinon.stub();
136149
const listener2 = sinon.stub();
137150

@@ -146,3 +159,105 @@ describe('BaseController', () => {
146159
expect(listener2.callCount).toEqual(0);
147160
});
148161
});
162+
163+
describe('getAnonymizedState', () => {
164+
it('should return empty state', () => {
165+
expect(getAnonymizedState({}, {})).toEqual({});
166+
});
167+
168+
it('should return empty state when no properties are anonymized', () => {
169+
const anonymizedState = getAnonymizedState({ count: 1 }, { count: { anonymous: false, persist: false } });
170+
expect(anonymizedState).toEqual({});
171+
});
172+
173+
it('should return state that is already anonymized', () => {
174+
const anonymizedState = getAnonymizedState(
175+
{
176+
password: 'secret password',
177+
privateKey: '123',
178+
network: 'mainnet',
179+
tokens: ['DAI', 'USDC'],
180+
},
181+
{
182+
password: {
183+
anonymous: false,
184+
persist: false,
185+
},
186+
privateKey: {
187+
anonymous: false,
188+
persist: false,
189+
},
190+
network: {
191+
anonymous: true,
192+
persist: false,
193+
},
194+
tokens: {
195+
anonymous: true,
196+
persist: false,
197+
},
198+
},
199+
);
200+
expect(anonymizedState).toEqual({ network: 'mainnet', tokens: ['DAI', 'USDC'] });
201+
});
202+
203+
it('should use anonymizing function to anonymize state', () => {
204+
const anonymizeTransactionHash = (hash: string) => {
205+
return hash.split('').reverse().join('');
206+
};
207+
208+
const anonymizedState = getAnonymizedState(
209+
{
210+
transactionHash: '0x1234',
211+
},
212+
{
213+
transactionHash: {
214+
anonymous: anonymizeTransactionHash,
215+
persist: false,
216+
},
217+
},
218+
);
219+
220+
expect(anonymizedState).toEqual({ transactionHash: '4321x0' });
221+
});
222+
});
223+
224+
describe('getPersistentState', () => {
225+
it('should return empty state', () => {
226+
expect(getPersistentState({}, {})).toEqual({});
227+
});
228+
229+
it('should return empty state when no properties are persistent', () => {
230+
const persistentState = getPersistentState({ count: 1 }, { count: { anonymous: false, persist: false } });
231+
expect(persistentState).toEqual({});
232+
});
233+
234+
it('should return persistent state', () => {
235+
const persistentState = getPersistentState(
236+
{
237+
password: 'secret password',
238+
privateKey: '123',
239+
network: 'mainnet',
240+
tokens: ['DAI', 'USDC'],
241+
},
242+
{
243+
password: {
244+
anonymous: false,
245+
persist: true,
246+
},
247+
privateKey: {
248+
anonymous: false,
249+
persist: true,
250+
},
251+
network: {
252+
anonymous: false,
253+
persist: false,
254+
},
255+
tokens: {
256+
anonymous: false,
257+
persist: false,
258+
},
259+
},
260+
);
261+
expect(persistentState).toEqual({ password: 'secret password', privateKey: '123' });
262+
});
263+
});

src/BaseControllerV2.ts

+43-1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,15 @@ enablePatches();
1111
*/
1212
export type Listener<T> = (state: T, patches: Patch[]) => void;
1313

14+
export type Anonymizer<T> = (value: T) => T;
15+
16+
export type Schema<T> = {
17+
[P in keyof T]: {
18+
persist: boolean;
19+
anonymous: boolean | Anonymizer<T[P]>;
20+
};
21+
};
22+
1423
/**
1524
* Controller class that provides state management and subscriptions
1625
*/
@@ -19,13 +28,18 @@ export class BaseController<S extends Record<string, any>> {
1928

2029
private internalListeners: Set<Listener<S>> = new Set();
2130

31+
public readonly schema: Schema<S>;
32+
2233
/**
2334
* Creates a BaseController instance.
2435
*
2536
* @param state - Initial controller state
37+
* @param schema - State schema, describing how to "anonymize" the state,
38+
* and which parts should be persisted.
2639
*/
27-
constructor(state: S) {
40+
constructor(state: S, schema: Schema<S>) {
2841
this.internalState = state;
42+
this.schema = schema;
2943
}
3044

3145
/**
@@ -89,3 +103,31 @@ export class BaseController<S extends Record<string, any>> {
89103
this.internalListeners.clear();
90104
}
91105
}
106+
107+
// This function acts as a type guard. Using a `typeof` conditional didn't seem to work.
108+
function isAnonymizingFunction<T>(x: boolean | Anonymizer<T>): x is Anonymizer<T> {
109+
return typeof x === 'function';
110+
}
111+
112+
export function getAnonymizedState<S extends Record<string, any>>(state: S, schema: Schema<S>) {
113+
return Object.keys(state).reduce((anonymizedState, _key) => {
114+
const key: keyof S = _key; // https://stackoverflow.com/questions/63893394/string-cannot-be-used-to-index-type-t
115+
const schemaValue = schema[key].anonymous;
116+
if (isAnonymizingFunction(schemaValue)) {
117+
anonymizedState[key] = schemaValue(state[key]);
118+
} else if (schemaValue) {
119+
anonymizedState[key] = state[key];
120+
}
121+
return anonymizedState;
122+
}, {} as Partial<S>);
123+
}
124+
125+
export function getPersistentState<S extends Record<string, any>>(state: S, schema: Schema<S>) {
126+
return Object.keys(state).reduce((persistedState, _key) => {
127+
const key: keyof S = _key; // https://stackoverflow.com/questions/63893394/string-cannot-be-used-to-index-type-t
128+
if (schema[key].persist) {
129+
persistedState[key] = state[key];
130+
}
131+
return persistedState;
132+
}, {} as Partial<S>);
133+
}

0 commit comments

Comments
 (0)