Skip to content

Commit 8479cc0

Browse files
lucy66hwlucy66hw
committed
more ut
Signed-off-by: xil <fridalu66@gmail.com>
1 parent f92c3ed commit 8479cc0

2 files changed

Lines changed: 175 additions & 33 deletions

File tree

‎tools/proto-convert/src/postprocessing/CleanupUnusedMessages.ts‎

Lines changed: 45 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -120,45 +120,30 @@ export function extractRootsFromServices(servicePath: string): string[] {
120120
return Array.from(roots);
121121
}
122122

123-
// ==================== CLI ====================
124-
125-
if (require.main === module) {
126-
const command = new Command()
127-
.description('Remove unused messages and enums from a proto file.')
128-
.addOption(new Option('-i, --input <path>', 'input proto file').default('protos/generated/models/aggregated_models.proto'))
129-
.addOption(new Option('-o, --output <path>', 'output proto file (defaults to input)'))
130-
.addOption(new Option('-s, --service <path>', 'service proto file to auto-detect roots')
131-
.default('protos/generated/services/default_service.proto'))
132-
.addOption(new Option('-r, --roots <names>', 'root message names (comma-separated, overrides --service)')
133-
.argParser((val: string) => val.split(',').map(s => s.trim())))
134-
.allowExcessArguments(false)
135-
.parse();
136-
137-
type CleanupOpts = {
138-
input: string;
139-
output?: string;
140-
service: string;
141-
roots?: string[];
142-
};
143-
144-
const opts = command.opts() as CleanupOpts;
123+
export type CleanupOptions = {
124+
input: string;
125+
output?: string;
126+
service?: string;
127+
roots?: string[];
128+
};
145129

130+
/**
131+
* Clean up unused messages and enums from a proto file.
132+
* Returns the number of removed messages and enums.
133+
*/
134+
export function cleanupUnusedMessages(opts: CleanupOptions): { removedMessages: number; removedEnums: number } {
146135
if (!existsSync(opts.input)) {
147-
logger.error(`Input file not found: ${opts.input}`);
148-
process.exit(1);
136+
throw new Error(`Input file not found: ${opts.input}`);
149137
}
150138

151139
// Get roots
152140
let roots: string[];
153141
if (opts.roots && opts.roots.length > 0) {
154142
roots = opts.roots;
155-
logger.info(`Using manually specified roots: ${roots.join(', ')}`);
156-
} else if (existsSync(opts.service)) {
143+
} else if (opts.service && existsSync(opts.service)) {
157144
roots = extractRootsFromServices(opts.service);
158-
logger.info(`Auto-detected roots from ${opts.service}: ${roots.join(', ')}`);
159145
} else {
160-
logger.error(`Service file not found: ${opts.service}. Specify --roots manually.`);
161-
process.exit(1);
146+
throw new Error(`Service file not found: ${opts.service}. Specify roots manually.`);
162147
}
163148

164149
const parsed = parseProtoFile(opts.input);
@@ -167,8 +152,7 @@ if (require.main === module) {
167152
const messageNames = new Set(parsed.messages.map(m => m.name));
168153
for (const rootMsg of roots) {
169154
if (!messageNames.has(rootMsg)) {
170-
logger.error(`Root message not found: ${rootMsg}`);
171-
process.exit(1);
155+
throw new Error(`Root message not found: ${rootMsg}`);
172156
}
173157
}
174158

@@ -179,8 +163,37 @@ if (require.main === module) {
179163
const keptMessages = filterMessages(parsed.messages, reachable);
180164
const keptEnums = filterEnums(parsed.enums, reachable);
181165

166+
const removedMessages = parsed.messages.length - keptMessages.length;
167+
const removedEnums = parsed.enums.length - keptEnums.length;
168+
182169
// Write output
183170
const outputPath = opts.output || opts.input;
184171
writeProtoFile(keptMessages, keptEnums, outputPath);
185-
logger.info(`Updated: ${outputPath}`);
172+
173+
return { removedMessages, removedEnums };
174+
}
175+
176+
// ==================== CLI ====================
177+
178+
if (require.main === module) {
179+
const command = new Command()
180+
.description('Remove unused messages and enums from a proto file.')
181+
.addOption(new Option('-i, --input <path>', 'input proto file').default('protos/generated/models/aggregated_models.proto'))
182+
.addOption(new Option('-o, --output <path>', 'output proto file (defaults to input)'))
183+
.addOption(new Option('-s, --service <path>', 'service proto file to auto-detect roots')
184+
.default('protos/generated/services/default_service.proto'))
185+
.addOption(new Option('-r, --roots <names>', 'root message names (comma-separated, overrides --service)')
186+
.argParser((val: string) => val.split(',').map(s => s.trim())))
187+
.allowExcessArguments(false)
188+
.parse();
189+
190+
const opts = command.opts() as CleanupOptions;
191+
192+
try {
193+
const { removedMessages, removedEnums } = cleanupUnusedMessages(opts);
194+
logger.info(`Removed ${removedMessages} messages, ${removedEnums} enums. Updated: ${opts.output || opts.input}`);
195+
} catch (error) {
196+
logger.error((error as Error).message);
197+
process.exit(1);
198+
}
186199
}

‎tools/proto-convert/test/postprocessing/CleanupUnusedMessages.test.ts‎

Lines changed: 130 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,32 @@
33
*/
44

55
import * as path from 'path';
6+
import * as fs from 'fs';
7+
import * as os from 'os';
68
import { parseProtoFile } from '../../src/postprocessing/parser';
79
import {
810
isBuiltInType,
911
findReachableTypes,
1012
filterMessages,
1113
filterEnums,
12-
extractRootsFromServices
14+
extractRootsFromServices,
15+
cleanupUnusedMessages,
16+
CleanupOptions
1317
} from '../../src/postprocessing/CleanupUnusedMessages';
1418
import { ProtoMessage, ProtoEnum } from '../../src/postprocessing/types';
1519

1620
const TEST_PROTO = path.join(__dirname, '../fixtures/proto/test.proto');
1721
const TEST_SERVICE_PROTO = path.join(__dirname, '../fixtures/proto/test_service.proto');
1822

23+
// Proto content for tests
24+
const PROTO_SERVICE_WITH_ROOTS = `
25+
syntax = "proto3";
26+
package test;
27+
28+
message SearchRequest { string query = 1; }
29+
message SearchResponse { string result = 1; }
30+
`;
31+
1932
describe('CleanupUnusedMessages', () => {
2033
const parsed = parseProtoFile(TEST_PROTO);
2134

@@ -267,3 +280,119 @@ describe('isBuiltInType', () => {
267280
expect(isBuiltInType('map<int32, CustomType>')).toBe(false);
268281
});
269282
});
283+
284+
describe('cleanupUnusedMessages', () => {
285+
let tempDir: string;
286+
let outputPath: string;
287+
288+
beforeEach(() => {
289+
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'cleanup-test-'));
290+
outputPath = path.join(tempDir, 'output.proto');
291+
});
292+
293+
afterEach(() => {
294+
if (fs.existsSync(tempDir)) {
295+
fs.rmSync(tempDir, { recursive: true });
296+
}
297+
});
298+
299+
it('should cleanup unused messages with manual roots', () => {
300+
const opts: CleanupOptions = {
301+
input: TEST_PROTO,
302+
output: outputPath,
303+
roots: ['SearchRequest', 'SearchResponse']
304+
};
305+
306+
const result = cleanupUnusedMessages(opts);
307+
308+
expect(result.removedMessages).toBeGreaterThan(0);
309+
expect(fs.existsSync(outputPath)).toBe(true);
310+
311+
const output = parseProtoFile(outputPath);
312+
const messageNames = output.messages.map(m => m.name);
313+
314+
expect(messageNames).toContain('SearchRequest');
315+
expect(messageNames).toContain('SearchResponse');
316+
expect(messageNames).not.toContain('UnusedMessage');
317+
});
318+
319+
it('should cleanup unused messages with service file', () => {
320+
const testService = path.join(tempDir, 'test_service.proto');
321+
fs.writeFileSync(testService, PROTO_SERVICE_WITH_ROOTS);
322+
323+
const opts: CleanupOptions = {
324+
input: TEST_PROTO,
325+
output: outputPath,
326+
service: testService
327+
};
328+
329+
const result = cleanupUnusedMessages(opts);
330+
331+
expect(fs.existsSync(outputPath)).toBe(true);
332+
expect(result.removedMessages).toBeGreaterThanOrEqual(0);
333+
});
334+
335+
it('should throw error if input file not found', () => {
336+
const opts: CleanupOptions = {
337+
input: '/non/existent/path.proto',
338+
output: outputPath,
339+
roots: ['SomeMessage']
340+
};
341+
342+
expect(() => cleanupUnusedMessages(opts)).toThrow('Input file not found');
343+
});
344+
345+
it('should throw error if service file not found and no roots specified', () => {
346+
const opts: CleanupOptions = {
347+
input: TEST_PROTO,
348+
output: outputPath,
349+
service: '/non/existent/service.proto'
350+
};
351+
352+
expect(() => cleanupUnusedMessages(opts)).toThrow('Service file not found');
353+
});
354+
355+
it('should throw error if root message not found', () => {
356+
const opts: CleanupOptions = {
357+
input: TEST_PROTO,
358+
output: outputPath,
359+
roots: ['NonExistentMessage']
360+
};
361+
362+
expect(() => cleanupUnusedMessages(opts)).toThrow('Root message not found');
363+
});
364+
365+
it('should write to input file if no output specified', () => {
366+
// Copy test proto to temp location
367+
const tempInput = path.join(tempDir, 'input.proto');
368+
fs.copyFileSync(TEST_PROTO, tempInput);
369+
370+
const opts: CleanupOptions = {
371+
input: tempInput,
372+
roots: ['SearchRequest', 'SearchResponse']
373+
};
374+
375+
cleanupUnusedMessages(opts);
376+
377+
// Should have written to input file
378+
expect(fs.existsSync(tempInput)).toBe(true);
379+
380+
const output = parseProtoFile(tempInput);
381+
expect(output.messages.map(m => m.name)).not.toContain('UnusedMessage');
382+
});
383+
384+
it('should return count of removed messages and enums', () => {
385+
const opts: CleanupOptions = {
386+
input: TEST_PROTO,
387+
output: outputPath,
388+
roots: ['SearchRequest', 'SearchResponse']
389+
};
390+
391+
const result = cleanupUnusedMessages(opts);
392+
393+
expect(typeof result.removedMessages).toBe('number');
394+
expect(typeof result.removedEnums).toBe('number');
395+
expect(result.removedMessages).toBeGreaterThanOrEqual(0);
396+
expect(result.removedEnums).toBeGreaterThanOrEqual(0);
397+
});
398+
});

0 commit comments

Comments
 (0)