Throw an error if webhook URL is invalid (replicate#84)

mattt · web-flow · commit 7ed6bb104793 · 2023-06-12T03:14:37.000-07:00
diff --git a/index.test.ts b/index.test.ts
@@ -151,6 +151,18 @@ describe('Replicate client', () => {
       });
       expect(prediction.id).toBe('ufawqhfynnddngldkgtslldrkq');
     });
+
+    test('Throws an error if webhook URL is invalid', async () => {
+      await expect(async () => {
+        await client.predictions.create({
+          version: '5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa',
+          input: {
+            text: 'Alice',
+          },
+          webhook: 'invalid-url',
+        });
+      }).rejects.toThrow('Invalid webhook URL');
+    });
     // Add more tests for error handling, edge cases, etc.
   });
 
@@ -321,6 +333,23 @@ describe('Replicate client', () => {
       expect(training.id).toBe('zz4ibbonubfz7carwiefibzgga');
     });
 
+    test('Throws an error if webhook is not a valid URL', async () => {
+      await expect(
+        client.trainings.create(
+          'owner',
+          'model',
+          '632231d0d49d34d5c4633bd838aee3d81d936e59a886fbf28524702003b4c532',
+          {
+            destination: 'new_owner/new_model',
+            input: {
+              text: '...',
+            },
+            webhook: 'invalid-url',
+          }
+        )
+      ).rejects.toThrow('Invalid webhook URL');
+    });
+
     // Add more tests for error handling, edge cases, etc.
   });
 
@@ -499,6 +528,18 @@ describe('Replicate client', () => {
       // @ts-expect-error
       await expect(client.run(':abc123', options)).rejects.toThrow();
     });
+
+    test('Throws an error if webhook URL is invalid', async () => {
+      await expect(async () => {
+        await client.run(
+          'owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa', {
+          input: {
+            text: 'Alice',
+          },
+          webhook: 'invalid-url',
+        });
+      }).rejects.toThrow('Invalid webhook URL');
+    });
   });
 
   // Continue with tests for other methods
diff --git a/lib/predictions.js b/lib/predictions.js
@@ -14,6 +14,15 @@
 async function createPrediction(options) {
   const { wait, ...data } = options;
 
+  if (data.webhook) {
+    try {
+      // eslint-disable-next-line no-new
+      new URL(data.webhook);
+    } catch (err) {
+      throw new Error('Invalid webhook URL');
+    }
+  }
+
   const prediction = this.request('/predictions', {
     method: 'POST',
     data,
diff --git a/lib/trainings.js b/lib/trainings.js
@@ -14,6 +14,15 @@
 async function createTraining(model_owner, model_name, version_id, options) {
   const { ...data } = options;
 
+  if (data.webhook) {
+    try {
+      // eslint-disable-next-line no-new
+      new URL(data.webhook);
+    } catch (err) {
+      throw new Error('Invalid webhook URL');
+    }
+  }
+
   const training = this.request(`/models/${model_owner}/${model_name}/versions/${version_id}/trainings`, {
     method: 'POST',
     data,
