@od-oneapp/ai-platform 0.1.0

package/src/output/__tests__/output.test.ts

@@ -0,0 +1,737 @@
+/**
+ * @fileoverview Output Module Tests
+ *
+ * Tests for createOutputValidator and withElementStream.
+ */
+
+import { describe, it, expect, vi } from 'vitest';
+import { z, ZodError } from 'zod';
+
+import {
+  withElementStream,
+  collectFromElementStream,
+  type ExperimentalOutputStreamResult,
+} from '../element-stream';
+import {
+  createOutputValidator,
+  createStrictValidator,
+  createPartialValidator,
+  isOutputValidationError,
+  OutputValidationError,
+  type PartialValidationResult,
+} from '../validator';
+
+// -----------------------------------------------------------------------------
+// Test Schemas
+// -----------------------------------------------------------------------------
+
+const PersonSchema = z.object({
+  name: z.string(),
+  age: z.number().min(0).max(150),
+  email: z.string().email(),
+});
+
+type Person = z.infer<typeof PersonSchema>;
+
+type Task = {
+  id: number;
+  title: string;
+  done: boolean;
+};
+
+// -----------------------------------------------------------------------------
+// createOutputValidator Tests
+// -----------------------------------------------------------------------------
+
+describe('createOutputValidator', () => {
+  describe('throw strategy', () => {
+    it('returns validated data for valid input', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'throw',
+      });
+
+      const input = { name: 'John', age: 30, email: 'john@example.com' };
+      const result = validator.validate(input);
+
+      expect(result).toStrictEqual(input);
+    });
+
+    it('throws OutputValidationError for invalid input', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'throw',
+      });
+
+      const input = { name: 'John', age: 'thirty', email: 'not-an-email' };
+
+      expect(() => validator.validate(input)).toThrowError(OutputValidationError);
+    });
+
+    it('includes error details in thrown error', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'throw',
+      });
+
+      const input = { name: 'John', age: -5, email: 'invalid' };
+
+      try {
+        validator.validate(input);
+        expect.fail('Should have thrown');
+      } catch (error) {
+        expect(isOutputValidationError(error)).toBe(true);
+        if (isOutputValidationError(error)) {
+          expect(error.zodError.issues.length).toBeGreaterThan(0);
+          expect(error.strategy).toBe('throw');
+          expect(error.attempts).toBe(1);
+        }
+      }
+    });
+
+    it('uses custom error message prefix', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'throw',
+        errorMessage: 'Person validation failed',
+      });
+
+      const input = { name: 123, age: 30, email: 'john@example.com' };
+
+      try {
+        validator.validate(input);
+        expect.fail('Should have thrown');
+      } catch (error) {
+        expect(error).toBeInstanceOf(OutputValidationError);
+        expect((error as OutputValidationError).message).toContain('Person validation failed');
+      }
+    });
+
+    it('calls onError callback on validation failure', () => {
+      const onError = vi.fn();
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'throw',
+        onError,
+      });
+
+      const input = { name: 'John', age: 200, email: 'john@example.com' };
+
+      try {
+        validator.validate(input);
+      } catch {
+        // Expected
+      }
+
+      expect(onError).toHaveBeenCalledTimes(1);
+      expect(onError).toHaveBeenCalledWith(expect.any(Object), 1);
+    });
+  });
+
+  describe('partial strategy', () => {
+    it('returns success result for valid input', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'partial',
+      });
+
+      const input = { name: 'Jane', age: 25, email: 'jane@example.com' };
+      const result = validator.validate(input) as PartialValidationResult<Person>;
+
+      // For partial strategy, validate() returns PartialValidationResult on success
+      // But we need to check how the implementation works
+      expect(result).toBeDefined();
+    });
+
+    it('returns partial data on validation failure', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'partial',
+      });
+
+      const input = { name: 'Jane', age: 25, email: 'not-valid' };
+      const result = validator.validate(input) as PartialValidationResult<Person>;
+
+      expect(result.success).toBe(false);
+      expect(result.partial).toBeDefined();
+      expect(result.partial?.name).toBe('Jane');
+      expect(result.partial?.age).toBe(25);
+      expect(result.errors).toBeDefined();
+      expect(result.errors!.length).toBeGreaterThan(0);
+    });
+
+    it('returns errors array with issue details', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'partial',
+      });
+
+      const input = { name: 'Jane', age: -10, email: 'bad' };
+      const result = validator.validate(input) as PartialValidationResult<Person>;
+
+      expect(result.success).toBe(false);
+      expect(result.errors).toBeDefined();
+      expect(result.errors!.some(e => e.path.includes('age'))).toBe(true);
+      expect(result.errors!.some(e => e.path.includes('email'))).toBe(true);
+    });
+
+    it('returns undefined partial for completely invalid input', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'partial',
+      });
+
+      const input = 'not an object';
+      const result = validator.validate(input) as PartialValidationResult<Person>;
+
+      expect(result.success).toBe(false);
+      expect(result.partial).toBeUndefined();
+    });
+  });
+
+  describe('retry strategy', () => {
+    it('throws when using sync validate with retry strategy', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'retry',
+        maxRetries: 2,
+      });
+
+      const input = { name: 'John', age: 'bad', email: 'invalid' };
+
+      expect(() => validator.validate(input)).toThrowError(OutputValidationError);
+    });
+
+    it('returns data on first attempt if valid', async () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'retry',
+        maxRetries: 2,
+      });
+
+      const input = { name: 'John', age: 30, email: 'john@example.com' };
+      const result = await validator.validateAsync(input);
+
+      expect(result).toStrictEqual(input);
+    });
+
+    it('throws if retry exhausted without valid result', async () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'retry',
+        maxRetries: 1,
+        retryDelay: 10,
+      });
+
+      const input = { name: 'John', age: 'bad', email: 'invalid' };
+
+      await expect(validator.validateAsync(input)).rejects.toThrowError(OutputValidationError);
+    });
+
+    it('throws when retryContext is missing for retry strategy', async () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'retry',
+        maxRetries: 1,
+      });
+
+      const input = { name: 'John', age: 'bad', email: 'invalid' };
+
+      await expect(validator.validateAsync(input)).rejects.toThrowError(/retryContext/);
+    });
+  });
+
+  describe('getSchema and getStrategy', () => {
+    it('returns the schema', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'throw',
+      });
+
+      expect(validator.getSchema()).toBe(PersonSchema);
+    });
+
+    it('returns the strategy', () => {
+      const validator = createOutputValidator({
+        schema: PersonSchema,
+        onValidationError: 'partial',
+      });
+
+      expect(validator.getStrategy()).toBe('partial');
+    });
+  });
+});
+
+// -----------------------------------------------------------------------------
+// createStrictValidator Tests
+// -----------------------------------------------------------------------------
+
+describe('createStrictValidator', () => {
+  it('returns validated data for valid input', () => {
+    const validate = createStrictValidator(PersonSchema);
+
+    const input = { name: 'Alice', age: 28, email: 'alice@example.com' };
+    const result = validate(input);
+
+    expect(result).toStrictEqual(input);
+  });
+
+  it('throws for invalid input', () => {
+    const validate = createStrictValidator(PersonSchema);
+
+    const input = { name: 'Alice', age: 'not-a-number', email: 'alice@example.com' };
+
+    expect(() => validate(input)).toThrowError(OutputValidationError);
+  });
+});
+
+// -----------------------------------------------------------------------------
+// createPartialValidator Tests
+// -----------------------------------------------------------------------------
+
+describe('createPartialValidator', () => {
+  it('returns success result for valid input', () => {
+    const validate = createPartialValidator(PersonSchema);
+
+    const input = { name: 'Bob', age: 35, email: 'bob@example.com' };
+    const result = validate(input);
+
+    expect(result.success).toBe(true);
+    expect(result.data).toStrictEqual(input);
+  });
+
+  it('returns partial result for invalid input', () => {
+    const validate = createPartialValidator(PersonSchema);
+
+    const input = { name: 'Bob', age: 35, email: 'invalid-email' };
+    const result = validate(input);
+
+    expect(result.success).toBe(false);
+    expect(result.partial?.name).toBe('Bob');
+    expect(result.partial?.age).toBe(35);
+    expect(result.errors).toBeDefined();
+  });
+});
+
+// -----------------------------------------------------------------------------
+// isOutputValidationError Tests
+// -----------------------------------------------------------------------------
+
+describe('isOutputValidationError', () => {
+  it('returns true for OutputValidationError', () => {
+    const error = new OutputValidationError('test', new ZodError([]), 'throw', 1);
+
+    expect(isOutputValidationError(error)).toBe(true);
+  });
+
+  it('returns false for regular Error', () => {
+    const error = new Error('regular error');
+
+    expect(isOutputValidationError(error)).toBe(false);
+  });
+
+  it('returns false for non-Error values', () => {
+    expect(isOutputValidationError('string')).toBe(false);
+    expect(isOutputValidationError(null)).toBe(false);
+    expect(isOutputValidationError(undefined)).toBe(false);
+    expect(isOutputValidationError({})).toBe(false);
+  });
+});
+
+// -----------------------------------------------------------------------------
+// withElementStream Tests
+// -----------------------------------------------------------------------------
+
+describe('withElementStream', () => {
+  // Helper to create mock stream result
+  function createMockStreamResult<T>(elements: T[]): ExperimentalOutputStreamResult<T> {
+    return {
+      experimental_output: {
+        elementStream: (async function* () {
+          for (const element of elements) {
+            yield element;
+          }
+        })(),
+        object: Promise.resolve(elements),
+      },
+    };
+  }
+
+  // Helper to create mock stream result with only object promise (no elementStream)
+  function createMockObjectResult<T>(elements: T[]): ExperimentalOutputStreamResult<T> {
+    return {
+      experimental_output: {
+        object: Promise.resolve(elements),
+      },
+    };
+  }
+
+  describe('basic iteration', () => {
+    it('yields elements from stream', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'Task 1', done: false },
+        { id: 2, title: 'Task 2', done: true },
+        { id: 3, title: 'Task 3', done: false },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const collected: Task[] = [];
+
+      for await (const task of withElementStream(streamResult)) {
+        collected.push(task as Task);
+      }
+
+      expect(collected).toHaveLength(3);
+      expect(collected).toStrictEqual(tasks);
+    });
+
+    it('handles empty stream', async () => {
+      const streamResult = createMockStreamResult<Task>([]);
+      const collected: Task[] = [];
+
+      for await (const task of withElementStream(streamResult)) {
+        collected.push(task as Task);
+      }
+
+      expect(collected).toHaveLength(0);
+    });
+
+    it('handles missing experimental_output', async () => {
+      const streamResult: ExperimentalOutputStreamResult<Task> = {};
+      const collected: Task[] = [];
+
+      for await (const task of withElementStream(streamResult)) {
+        collected.push(task as Task);
+      }
+
+      expect(collected).toHaveLength(0);
+    });
+  });
+
+  describe('onElement callback', () => {
+    it('calls onElement for each element', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'Task 1', done: false },
+        { id: 2, title: 'Task 2', done: true },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const onElement = vi.fn();
+
+      for await (const _ of withElementStream(streamResult, { onElement })) {
+        // Just iterate
+      }
+
+      expect(onElement).toHaveBeenCalledTimes(2);
+      expect(onElement).toHaveBeenNthCalledWith(1, tasks[0], 0);
+      expect(onElement).toHaveBeenNthCalledWith(2, tasks[1], 1);
+    });
+
+    it('provides correct index to onElement', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'A', done: false },
+        { id: 2, title: 'B', done: false },
+        { id: 3, title: 'C', done: false },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const indices: number[] = [];
+
+      for await (const _ of withElementStream(streamResult, {
+        onElement: (_, index) => indices.push(index),
+      })) {
+        // Just iterate
+      }
+
+      expect(indices).toStrictEqual([0, 1, 2]);
+    });
+  });
+
+  describe('transform', () => {
+    it('transforms elements before yielding', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'task', done: false },
+        { id: 2, title: 'another', done: true },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const collected: Task[] = [];
+
+      for await (const task of withElementStream(streamResult, {
+        transform: t => ({ ...t, title: t.title.toUpperCase() }),
+      })) {
+        collected.push(task as Task);
+      }
+
+      expect(collected[0]!.title).toBe('TASK');
+      expect(collected[1]!.title).toBe('ANOTHER');
+    });
+
+    it('calls onElement with transformed elements', async () => {
+      const tasks: Task[] = [{ id: 1, title: 'test', done: false }];
+
+      const streamResult = createMockStreamResult(tasks);
+      const onElement = vi.fn();
+
+      for await (const _ of withElementStream(streamResult, {
+        transform: t => ({ ...t, title: 'TRANSFORMED' }),
+        onElement,
+      })) {
+        // Just iterate
+      }
+
+      expect(onElement).toHaveBeenCalledWith(expect.objectContaining({ title: 'TRANSFORMED' }), 0);
+    });
+  });
+
+  describe('bufferSize', () => {
+    it('buffers elements when bufferSize > 1', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'A', done: false },
+        { id: 2, title: 'B', done: false },
+        { id: 3, title: 'C', done: false },
+        { id: 4, title: 'D', done: false },
+        { id: 5, title: 'E', done: false },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const batches: Task[][] = [];
+
+      for await (const batch of withElementStream(streamResult, {
+        bufferSize: 2,
+      })) {
+        batches.push(batch as Task[]);
+      }
+
+      expect(batches).toHaveLength(3); // 2, 2, 1
+      expect(batches[0]).toHaveLength(2);
+      expect(batches[1]).toHaveLength(2);
+      expect(batches[2]).toHaveLength(1);
+    });
+
+    it('still calls onElement for each individual element in buffered mode', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'A', done: false },
+        { id: 2, title: 'B', done: false },
+        { id: 3, title: 'C', done: false },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const onElement = vi.fn();
+
+      for await (const _ of withElementStream(streamResult, {
+        bufferSize: 2,
+        onElement,
+      })) {
+        // Just iterate
+      }
+
+      expect(onElement).toHaveBeenCalledTimes(3);
+    });
+  });
+
+  describe('onComplete', () => {
+    it('calls onComplete when stream finishes', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'A', done: false },
+        { id: 2, title: 'B', done: true },
+      ];
+
+      const streamResult = createMockStreamResult(tasks);
+      const onComplete = vi.fn();
+
+      for await (const _ of withElementStream(streamResult, { onComplete })) {
+        // Just iterate
+      }
+
+      expect(onComplete).toHaveBeenCalledTimes(1);
+      expect(onComplete).toHaveBeenCalledWith(tasks, 2);
+    });
+
+    it('calls onComplete with transformed elements', async () => {
+      const tasks: Task[] = [{ id: 1, title: 'test', done: false }];
+
+      const streamResult = createMockStreamResult(tasks);
+      const onComplete = vi.fn();
+
+      for await (const _ of withElementStream(streamResult, {
+        transform: t => ({ ...t, title: 'UPPER' }),
+        onComplete,
+      })) {
+        // Just iterate
+      }
+
+      expect(onComplete).toHaveBeenCalledWith([expect.objectContaining({ title: 'UPPER' })], 1);
+    });
+  });
+
+  describe('onError', () => {
+    it('calls onError when stream throws', async () => {
+      const error = new Error('Stream error');
+      const streamResult: ExperimentalOutputStreamResult<Task> = {
+        experimental_output: {
+          elementStream: (async function* () {
+            yield { id: 1, title: 'A', done: false };
+            throw error;
+          })(),
+        },
+      };
+
+      const onError = vi.fn();
+      const collected: Task[] = [];
+
+      try {
+        for await (const task of withElementStream(streamResult, { onError })) {
+          collected.push(task as Task);
+        }
+      } catch {
+        // Expected
+      }
+
+      expect(onError).toHaveBeenCalledTimes(1);
+      expect(onError).toHaveBeenCalledWith(error);
+    });
+
+    it('re-throws error after calling onError', async () => {
+      const error = new Error('Test error');
+      const streamResult: ExperimentalOutputStreamResult<Task> = {
+        experimental_output: {
+          elementStream: (async function* () {
+            throw error;
+          })(),
+        },
+      };
+
+      const onError = vi.fn();
+
+      await expect(async () => {
+        for await (const _ of withElementStream(streamResult, { onError })) {
+          // Should not reach here
+        }
+      }).rejects.toThrowError('Test error');
+
+      expect(onError).toHaveBeenCalledWith();
+    });
+  });
+
+  describe('fallback to object promise', () => {
+    it('uses object promise when elementStream is not available', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'Task 1', done: false },
+        { id: 2, title: 'Task 2', done: true },
+      ];
+
+      const streamResult = createMockObjectResult(tasks);
+      const collected: Task[] = [];
+
+      for await (const task of withElementStream(streamResult)) {
+        collected.push(task as Task);
+      }
+
+      expect(collected).toStrictEqual(tasks);
+    });
+
+    it('applies transform when using object promise', async () => {
+      const tasks: Task[] = [{ id: 1, title: 'lower', done: false }];
+
+      const streamResult = createMockObjectResult(tasks);
+      const collected: Task[] = [];
+
+      for await (const task of withElementStream(streamResult, {
+        transform: t => ({ ...t, title: t.title.toUpperCase() }),
+      })) {
+        collected.push(task as Task);
+      }
+
+      expect(collected[0]!.title).toBe('LOWER');
+    });
+
+    it('supports buffering when using object promise', async () => {
+      const tasks: Task[] = [
+        { id: 1, title: 'A', done: false },
+        { id: 2, title: 'B', done: false },
+        { id: 3, title: 'C', done: false },
+      ];
+
+      const streamResult = createMockObjectResult(tasks);
+      const batches: Task[][] = [];
+
+      for await (const batch of withElementStream(streamResult, {
+        bufferSize: 2,
+      })) {
+        batches.push(batch as Task[]);
+      }
+
+      expect(batches).toHaveLength(2);
+      expect(batches[0]).toHaveLength(2);
+      expect(batches[1]).toHaveLength(1);
+    });
+  });
+});
+
+// -----------------------------------------------------------------------------
+// collectFromElementStream Tests
+// -----------------------------------------------------------------------------
+
+describe('collectFromElementStream', () => {
+  function createMockStreamResult<T>(elements: T[]): ExperimentalOutputStreamResult<T> {
+    return {
+      experimental_output: {
+        elementStream: (async function* () {
+          for (const element of elements) {
+            yield element;
+          }
+        })(),
+      },
+    };
+  }
+
+  it('collects all elements into an array', async () => {
+    const tasks: Task[] = [
+      { id: 1, title: 'Task 1', done: false },
+      { id: 2, title: 'Task 2', done: true },
+    ];
+
+    const streamResult = createMockStreamResult(tasks);
+    const collected = await collectFromElementStream(streamResult);
+
+    expect(collected).toStrictEqual(tasks);
+  });
+
+  it('applies transform to collected elements', async () => {
+    const tasks: Task[] = [{ id: 1, title: 'lower', done: false }];
+
+    const streamResult = createMockStreamResult(tasks);
+    const collected = await collectFromElementStream(streamResult, {
+      transform: t => ({ ...t, title: t.title.toUpperCase() }),
+    });
+
+    expect(collected[0]!.title).toBe('LOWER');
+  });
+
+  it('calls onElement for each element', async () => {
+    const tasks: Task[] = [
+      { id: 1, title: 'A', done: false },
+      { id: 2, title: 'B', done: true },
+    ];
+
+    const streamResult = createMockStreamResult(tasks);
+    const onElement = vi.fn();
+
+    await collectFromElementStream(streamResult, { onElement });
+
+    expect(onElement).toHaveBeenCalledTimes(2);
+  });
+
+  it('calls onComplete when done', async () => {
+    const tasks: Task[] = [{ id: 1, title: 'Test', done: false }];
+
+    const streamResult = createMockStreamResult(tasks);
+    const onComplete = vi.fn();
+
+    await collectFromElementStream(streamResult, { onComplete });
+
+    expect(onComplete).toHaveBeenCalledTimes(1);
+    expect(onComplete).toHaveBeenCalledWith(tasks, 1);
+  });
+});