forked from AIObjectives/talk-to-the-city-reports
-
Notifications
You must be signed in to change notification settings - Fork 0
/
whisper_v1.test.ts
141 lines (132 loc) · 3.84 KB
/
whisper_v1.test.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import { describe, it, expect, vi, beforeEach } from 'vitest';
import WhisperNode, { whisper_node_data } from '$lib/compute/whisper_v1';
import deepCopy from 'deep-copy';
import { readFileFromGCS, uploadJSONToGCS } from '$lib/utils';
import MockOpenAI from '$lib/mock_open_ai';
import File from './test/mocks/File';
import sample from '$lib/mock_data/speech_to_text/sample.json';
globalThis.File = File;
vi.mock('$lib/utils', async () => {
const actual = await vi.importActual('$lib/utils');
return {
...actual,
readFileFromGCS: vi.fn(),
uploadJSONToGCS: vi.fn()
};
});
describe('WhisperNode class', () => {
let node;
let inputData;
let mockAudioFile;
let open_ai_key;
beforeEach(() => {
node = new WhisperNode(deepCopy(whisper_node_data));
node.data.response_format = 'verbose_json';
open_ai_key = 'test_open_ai_key';
mockAudioFile = new File(['audio content'], 'test.mp3', { type: 'audio/mpeg' });
inputData = {
audio: mockAudioFile,
open_ai_key: open_ai_key
};
vi.clearAllMocks();
});
it('should load from cache if data is not dirty and gcs_path is set', async () => {
node.data.dirty = false;
node.data.gcs_path = 'gs://bucket/path/test.json';
node.data.output = sample;
const output = await node.compute(
inputData,
'run',
vi.fn(),
vi.fn(),
vi.fn(),
'test_slug',
vi.fn()
);
expect(output).toEqual(sample);
});
it('should load from GCS if data is not dirty, gcs_path is set, and output is empty and audio size matches', async () => {
node.data.dirty = false;
node.data.gcs_path = 'gs://bucket/path/test.json';
node.data.output = null;
node.data.audio_size = mockAudioFile.size;
vi.mocked(readFileFromGCS).mockResolvedValue(JSON.stringify(sample));
const output = await node.compute(
inputData,
'run',
vi.fn(),
vi.fn(),
vi.fn(),
'test_slug',
vi.fn()
);
expect(readFileFromGCS).toHaveBeenCalled();
expect(output).toEqual(sample);
});
it('should transcribe audio and upload to GCS if data is dirty', async () => {
node.data.dirty = true;
node.data.gcs_path = 'gs://bucket/path/test.json';
vi.mocked(uploadJSONToGCS).mockResolvedValue(undefined);
MockOpenAI.prototype.audio = {
transcriptions: {
create: vi.fn().mockResolvedValue(sample)
}
};
const output = await node.compute(
inputData,
'run',
vi.fn(),
vi.fn(),
vi.fn(),
'test_slug',
vi.fn()
);
expect(uploadJSONToGCS).toHaveBeenCalled();
expect(output).toEqual(sample);
expect(node.data.dirty).toBe(false);
});
it('should return undefined and set message if open_ai_key is missing', async () => {
delete inputData.open_ai_key;
const output = await node.compute(
inputData,
'run',
vi.fn(),
vi.fn(),
vi.fn(),
'test_slug',
vi.fn()
);
expect(output).toBeUndefined();
expect(node.data.message).toContain('missing_input_data');
});
it('should convert transcription to internal format if response_format is custom', async () => {
node.data.dirty = true;
node.data.response_format = 'custom';
node.data.interview = 'bob';
node.data.video = 'https://vimeo.com/855112091';
vi.mocked(uploadJSONToGCS).mockResolvedValue(undefined);
MockOpenAI.prototype.audio = {
transcriptions: {
create: vi.fn().mockResolvedValue(sample)
}
};
const output = await node.compute(
inputData,
'run',
vi.fn(),
vi.fn(),
vi.fn(),
'test_slug',
vi.fn()
);
expect(output).toEqual([
{
'comment-id': 'whisper_0',
'comment-body': 'This is an audio sample.',
timestamp: '00:00:00',
interview: 'bob',
video: 'https://vimeo.com/855112091'
}
]);
});
});