Add tests for stream option in run operation

Author: mattt

index.test.ts

@@ -2,6 +2,21 @@ import { expect, jest, test } from '@jest/globals';
 import Replicate, { Prediction } from 'replicate';
 import nock from 'nock';
 import fetch from 'cross-fetch';
+import { SetupServer, setupServer } from 'msw/node';
+import { rest } from 'msw';
+
+type ServerOptions = Parameters<typeof setupServer>;
+
+const withRequestInterception =
+  (handlers: ServerOptions, test: (server: SetupServer) => any) => async () => {
+    const server = setupServer(...handlers);
+    server.listen();
+
+    return Promise.resolve(test(server)).finally(() => {
+      server.resetHandlers();
+      server.close();
+    });
+  };
 
 describe('Replicate client', () => {
   let client: Replicate;
@@ -483,7 +498,7 @@ describe('Replicate client', () => {
           status: 'processing',
         })
         .get('/predictions/ufawqhfynnddngldkgtslldrkq')
-        .reply(200, {
+        .reply(201, {
           id: 'ufawqhfynnddngldkgtslldrkq',
           status: 'succeeded',
           output: 'foobar',
@@ -498,6 +513,29 @@ describe('Replicate client', () => {
       expect(output).toBe('foobar');
     });
 
+    test('Sends an Accept: text/event-stream header when stream is true', async () => {
+      withRequestInterception(
+        [
+          rest.post('https://api.replicate.com/v1/predictions/', (req, res, ctx) => {
+            expect(req.headers.get('Accept')).toBe('text/event-stream');
+            return res(ctx.status(201), ctx.text('data: Once upon a time\n\n'));
+          }),
+        ],
+        async () => {
+          const stream = await client.run(
+            'owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa',
+            {
+              input: { prompt: 'Tell me a story' },
+              stream: true
+            }
+          );
+
+          for await (const event of await stream) {
+            expect(event).toBe('Once upon a time');
+          }
+        });
+    });
+
     test('Does not throw an error for identifier containing hyphen and full stop', async () => {
       nock(BASE_URL)
         .post('/predictions')
@@ -520,12 +558,10 @@ describe('Replicate client', () => {
 
       await expect(client.run('owner/model:invalid', options)).rejects.toThrow();
 
-      // @ts-expect-error
       await expect(client.run('owner:abc123', options)).rejects.toThrow();
 
       await expect(client.run('/model:abc123', options)).rejects.toThrow();
 
-      // @ts-expect-error
       await expect(client.run(':abc123', options)).rejects.toThrow();
     });