greatly improve schema and type checking

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@nothing-special/kaiware-lib",
-	"version": "0.4.0",
+	"version": "0.5.0",
 	"type": "module",
 	"author": {
 		"name": "Garrett Downs",

src/enums/MessageType.ts

@@ -7,22 +7,27 @@ export enum MessageType {
 
 	GetElementStyles = 'get-element-styles',
 	GetElementStylesRes = 'get-element-styles-res',
+
 	SetElementStyles = 'set-element-styles',
 	SetElementStylesRes = 'set-element-styles-res',
 
 	GetElementData = 'get-element-data',
 	GetElementDataRes = 'get-element-data-res',
+
 	SetElementData = 'set-element-data',
 	SetElementDataRes = 'set-element-data-res',
 
 	GetStorage = 'get-storage',
 	GetStorageRes = 'get-storage-res',
+
 	SetStorage = 'set-storage',
 	SetStorageRes = 'set-storage-res',
 
+	ClearLogs = 'clear-logs',
+	ClearLogsRes = 'clear-logs-res',
+
 	// Misc
 
 	NewLog = 'new-log',
-	ClearLogs = 'clear-logs',
 	Error = 'error'
 }

src/lib/connection.ts

@@ -1,20 +1,6 @@
 import { z } from 'zod';
 import { MessageType } from '../enums';
-import {
-	Config,
-	GetDeviceInfoPayload,
-	GetDeviceInfoResPayload,
-	GetElementDataPayload,
-	GetElementDataResPayload,
-	GetElementStylesPayload,
-	GetElementStylesResPayload,
-	GetElementsPayload,
-	GetElementsResPayload,
-	GetStoragePayload,
-	GetStorageResPayload,
-	Message
-} from '../types';
-import { isJson } from '../utils';
+import { Config, MessageWithId, rawMessageSchema } from '../types';
 
 export class Connection {
 	private socket: WebSocket | null = null;
@@ -41,18 +27,7 @@ export class Connection {
 			this.socket.onmessage = (event) => {
 				console.log('Message received: ', JSON.parse(event.data));
 
-				const messageSchema = z
-					.string()
-					.refine((val) => isJson(val), 'Must be a valid JSON string')
-					.transform((val) => JSON.parse(val))
-					.pipe(
-						z.object({
-							type: z.nativeEnum(MessageType),
-							data: z.any().optional()
-						})
-					);
-
-				const validateMessage = messageSchema.safeParse(event.data);
+				const validateMessage = rawMessageSchema.safeParse(event.data);
 				if (!validateMessage.success) {
 					console.log(
 						JSON.stringify(formatValidationError(validateMessage.error.issues))
@@ -61,7 +36,7 @@ export class Connection {
 				}
 
 				// eslint-disable-next-line @typescript-eslint/no-explicit-any
-				const message = JSON.parse(event.data) as Message<any>;
+				const message = JSON.parse(event.data) as MessageWithId;
 
 				switch (message.type) {
 					case MessageType.GetDeviceInfo:
@@ -105,7 +80,7 @@ export class Connection {
 		this.socket.close();
 	}
 
-	sendMessage<TData = undefined>(message: Message<TData>) {
+	sendMessage(message: MessageWithId) {
 		console.log('Sending message: ', message);
 
 		if (!this.socket) {
@@ -116,8 +91,8 @@ export class Connection {
 		this.socket.send(JSON.stringify(message));
 	}
 
-	private handleGetDeviceInfo(message: Message<GetDeviceInfoPayload>) {
-		this.sendMessage<GetDeviceInfoResPayload>({
+	private handleGetDeviceInfo(message: MessageWithId & { type: MessageType.GetDeviceInfo }) {
+		this.sendMessage({
 			requestId: message.requestId,
 			type: MessageType.GetDeviceInfoRes,
 			data: {
@@ -128,17 +103,17 @@ export class Connection {
 		});
 	}
 
-	private handleGetElements(message: Message<GetElementsPayload>) {
-		this.sendMessage<GetElementsResPayload>({
+	private handleGetElements(message: MessageWithId & { type: MessageType.GetElements }) {
+		this.sendMessage({
 			requestId: message.requestId,
 			type: MessageType.GetElementsRes,
 			data: document.querySelector('html')?.outerHTML ?? ''
 		});
 	}
 
-	private handleGetStorage(message: Message<GetStoragePayload>) {
+	private handleGetStorage(message: MessageWithId & { type: MessageType.GetStorage }) {
 		if (message.data.storageType === 'local') {
-			this.sendMessage<GetStorageResPayload>({
+			this.sendMessage({
 				requestId: message.requestId,
 				type: MessageType.GetStorageRes,
 				data: {
@@ -147,7 +122,7 @@ export class Connection {
 				}
 			});
 		} else if (message.data.storageType === 'session') {
-			this.sendMessage<GetStorageResPayload>({
+			this.sendMessage({
 				requestId: message.requestId,
 				type: MessageType.GetStorageRes,
 				data: {
@@ -158,7 +133,9 @@ export class Connection {
 		}
 	}
 
-	private handleGetElementStyles(message: Message<GetElementStylesPayload>) {
+	private handleGetElementStyles(
+		message: MessageWithId & { type: MessageType.GetElementStyles }
+	) {
 		console.log('Getting styles for element at index: ', message.data.index);
 
 		const element = document.querySelectorAll('*')[message.data.index];
@@ -172,7 +149,7 @@ export class Connection {
 			},
 			{} as Record<string, string>
 		);
-		this.sendMessage<GetElementStylesResPayload>({
+		this.sendMessage({
 			requestId: message.requestId,
 			type: MessageType.GetElementStylesRes,
 			data: {
@@ -182,18 +159,18 @@ export class Connection {
 		});
 	}
 
-	private handleGetElementData(message: Message<GetElementDataPayload>) {
+	private handleGetElementData(message: MessageWithId & { type: MessageType.GetElementData }) {
 		console.log('Getting data for element at index: ', message.data.index);
 
 		const element = document.querySelectorAll('*')[message.data.index];
 		console.log('Element: ', element);
 
 		// TODO: Get data for element
 
-		this.sendMessage<GetElementDataResPayload>({
+		this.sendMessage({
 			requestId: message.requestId,
 			type: MessageType.GetElementDataRes,
-			data: {}
+			data: { index: message.data.index, data: {} }
 		});
 	}
 }

src/lib/kaiware.ts

@@ -1,5 +1,5 @@
 import { LogLevel, MessageType } from '../enums';
-import { Config, GetLogResPayload } from '../types';
+import { Config } from '../types';
 import { parseError } from '../utils';
 import { Connection } from './connection';
 
@@ -52,7 +52,7 @@ export class Kaiware {
 			}
 		});
 
-		this.connection.sendMessage<GetLogResPayload>({
+		this.connection.sendMessage({
 			requestId: '',
 			type: MessageType.NewLog,
 			data: {

src/types/schemas.test.ts

@@ -0,0 +1,90 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { describe, expect, it } from 'vitest';
+import { LogLevel, MessageType } from '../enums';
+import { getDeviceInfoResPayloadSchema, messageSchema, rawMessageSchema } from './schemas';
+
+describe('schemas', () => {
+	describe('messages', () => {
+        const testCases: any[] = [
+            [{requestId:'req',type:MessageType.GetDeviceInfo, data: null}, true, ''],
+			[{requestId:'req',type:MessageType.GetDeviceInfoRes, data: { id: 'id', name: 'name', connectionType: 'wifi' }}, true, ''],
+			[{requestId:'req',type:MessageType.GetElements, data: null}, true, ''],
+			[{requestId:'req',type:MessageType.GetElementsRes, data: 'html'}, true, ''],
+			[{requestId:'req',type:MessageType.GetElementStyles, data:{ index: 1 }}, true, ''],
+			[{requestId:'req',type:MessageType.GetElementStylesRes, data:{ index: 1, styles: {color:'red'} }}, true, ''],
+			[{requestId:'req',type:MessageType.GetElementData, data:{ index: 1 }}, true, ''],
+			[{requestId:'req',type:MessageType.GetElementDataRes, data:{ index: 1, data: {height:'100px'} }}, true, ''],
+			[{requestId:'req',type:MessageType.GetStorage, data: {storageType: 'local'}}, true, ''],
+			[{requestId:'req',type:MessageType.GetStorageRes, data: {storageType: 'local', data: {}}}, true, ''],
+			[{requestId:'req',type:MessageType.SetStorage, data: {storageType: 'local', data: {}}}, true, ''],
+			[{requestId:'req',type:MessageType.SetStorageRes, data: null}, true, ''],
+			[{requestId:'req',type:MessageType.NewLog, data: {source: 'src', level: LogLevel.Info, data: ['1','2'], timestamp: '2024'}}, true, ''],
+			[{requestId:'req',type:MessageType.ClearLogs, data: null}, true, ''],
+			[{requestId:'req',type:MessageType.ClearLogsRes, data: null}, true, ''],
+            [{requestId: null, type:MessageType.GetDeviceInfo, data: null}, false, 'requestId'],
+			[{requestId:'req', type: null, data: null}, false, 'type'],
+			[{requestId:'req', type:MessageType.GetDeviceInfoRes, data: { id: 'id', name: null, connectionType: 'wifi' }}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetElements, data: 123}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetElementsRes, data: null}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetElementStyles, data:{ index: [] }}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetElementStylesRes, data:{ index: 1, styles: null }}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetElementData, data: null}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetElementDataRes, data:{ index: 'abc', data: {height:'100px'} }}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetStorage, data: {storageType: 'other'}}, false, 'data'],
+			[{requestId:'req', type:MessageType.GetStorageRes, data: {storageType: 'local', data: null}}, false, 'data'],
+			[{requestId:'req', type:MessageType.SetStorage, data: {storageType: null, data: {}}}, false, 'data'],
+			[{requestId:'req', type:MessageType.SetStorageRes, data: {prop:1}}, false, 'data'],
+			[{requestId:'req', type:MessageType.NewLog, data: {source: 'src', level: 'other', data: ['1','2'], timestamp: '2024'}}, false, 'data'],
+			[{requestId:'req', type:MessageType.ClearLogs, data: 123}, false, 'data'],
+			[{requestId:'req', type:MessageType.ClearLogsRes, data: 'abc'}, false, 'data']
+        ];
+
+		it.each(testCases)(`should parse %s and return %s`, (data, success, errMsg) => {
+			const result = messageSchema.safeParse(data);
+
+			expect(result.success).toBe(success);
+            if(success) {
+                expect((result as any).data).toEqual(data);
+            } else {
+                expect((result as any).error.issues[0].path[0]).toEqual(errMsg);
+            }
+		});
+
+		it.each(testCases.map(a => ([
+            JSON.stringify(a[0]),
+            a[1],
+            a[2]
+        ])))(`should parse %s and return %s`, (data, success, errMsg) => {
+			const result = rawMessageSchema.safeParse(data);
+
+			expect(result.success).toBe(success);
+            if(success) {
+                expect(JSON.stringify((result as any).data)).toEqual(data);
+            } else {
+                expect((result as any).error.issues[0].path[0]).toEqual(errMsg);
+            }
+		});
+	});
+
+    describe('getDeviceInfoResPayloadSchema', () => {
+		it.each([
+			[{ id: 'id', name: 'name', connectionType: 'wifi' }, true],
+			[{ id: 'id', name: 'name', connectionType: 'usb' }, true]
+		])(`should successfully parse %s`, (data, success) => {
+			const result = getDeviceInfoResPayloadSchema.safeParse(data);
+
+			expect(result.success).toBe(success);
+			expect((result as any).data).toEqual(data);
+		});
+
+		it.each([
+			[{ id: 'id', name: 'name', connectionType: 'other' }, false, 'connectionType'],
+			[{ id: undefined, name: 'name', connectionType: 'usb' }, false, 'id']
+		])(`should fail to parse %s`, (data, success, errMsg) => {
+			const result = getDeviceInfoResPayloadSchema.safeParse(data);
+
+			expect(result.success).toBe(success);
+			expect((result as any).error.issues[0].path[0]).toEqual(errMsg);
+		});
+	});
+});