Add some tests
This commit is contained in:
BIN
Binary file not shown.
@@ -1,6 +1,6 @@
|
|||||||
use clap::Parser;
|
use clap::Parser;
|
||||||
use roto::proto_gen::google::protobuf::descriptor::FileDescriptorSet;
|
use roto::proto_gen::google::protobuf::descriptor::FileDescriptorSet;
|
||||||
use roto::proto_gen::generator::generate_rust_code;
|
use roto::generator::generate_rust_code;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
|||||||
+65
-38
@@ -29,46 +29,49 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|||||||
|
|
||||||
let request = CodeGeneratorRequest::new(&stdin_buf)?;
|
let request = CodeGeneratorRequest::new(&stdin_buf)?;
|
||||||
|
|
||||||
|
// 2. Process request and get response bytes
|
||||||
|
let response_bytes = handle_request(&request)?;
|
||||||
|
|
||||||
|
// 3. Write response to stdout
|
||||||
|
io::stdout().write_all(&response_bytes)?;
|
||||||
|
|
||||||
|
info!("Successfully wrote CodeGeneratorResponse to stdout");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Core logic that transforms a CodeGeneratorRequest into a serialized CodeGeneratorResponse.
|
||||||
|
fn handle_request(request: &CodeGeneratorRequest) -> std::result::Result<Vec<u8>, Box<dyn std::error::Error>> {
|
||||||
// 2. Construct a FileDescriptorSet from the request's proto_files
|
// 2. Construct a FileDescriptorSet from the request's proto_files
|
||||||
// Since generate_rust_code expects a FileDescriptorSet, we wrap the proto_files in one.
|
let mut set_buf = Vec::new();
|
||||||
// A FileDescriptorSet is a message with a repeated field 'file' at tag 1.
|
for file_res in request.proto_file() {
|
||||||
let mut set_buf = vec![0u8; 1024 * 1024]; // Allocate 1MB for descriptors
|
let (file_data, _) = file_res.map_err(|e| {
|
||||||
{
|
error!("Failed to iterate proto_file: {:?}", e);
|
||||||
let mut builder = ProtoBuilder::new(&mut set_buf);
|
e
|
||||||
for file_res in request.proto_file() {
|
})?;
|
||||||
let (file_data, _) = file_res.map_err(|e| {
|
|
||||||
error!("Failed to iterate proto_file: {:?}", e);
|
// Tag 1, Length-delimited: (1 << 3) | 2 = 10
|
||||||
e
|
set_buf.push(10);
|
||||||
})?;
|
|
||||||
builder.write_bytes(1, file_data).map_err(|e| {
|
// Write length as varint
|
||||||
error!("Failed to write proto_file to set: {:?}", e);
|
let len = file_data.len() as u64;
|
||||||
e
|
let mut len_buf = [0u8; 10];
|
||||||
})?;
|
let len_size = roto::write_varint(len, &mut len_buf).map_err(|e| {
|
||||||
}
|
error!("Failed to write varint length: {:?}", e);
|
||||||
// We don't need to call finish() here as we just need the buffer to be populated.
|
e
|
||||||
|
})?;
|
||||||
|
set_buf.extend_from_slice(&len_buf[..len_size]);
|
||||||
|
|
||||||
|
// Write data
|
||||||
|
set_buf.extend_from_slice(file_data);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Trim set_buf to the actual size written by ProtoBuilder
|
|
||||||
// We need to find the actual length. ProtoBuilder doesn't expose pos publicly in the provided snippet,
|
|
||||||
// but we can wrap the result of the generation.
|
|
||||||
// Actually, since FileDescriptorSet::new just wraps the buffer, it will read until the end.
|
|
||||||
// We should be careful about trailing zeros.
|
|
||||||
// Let's just use a slice of the buffer that was actually used.
|
|
||||||
// Since ProtoBuilder's pos is private, we might need a different approach or just hope the
|
|
||||||
// ProtoAccessor handles trailing zeros (which it should if it's a valid proto message).
|
|
||||||
// However, the most robust way is to use the length.
|
|
||||||
|
|
||||||
// Let's try to use the buffer as is. If ProtoAccessor is correctly implemented,
|
|
||||||
// it will stop at the end of the message.
|
|
||||||
let set = FileDescriptorSet::new(&set_buf)?;
|
let set = FileDescriptorSet::new(&set_buf)?;
|
||||||
|
|
||||||
// 3. Generate the Rust code
|
// Generate the Rust code
|
||||||
info!("Generating Rust code from descriptor set...");
|
info!("Generating Rust code from descriptor set...");
|
||||||
let generated_code = generate_rust_code(&set);
|
let generated_code = generate_rust_code(&set);
|
||||||
|
|
||||||
// 4. Construct CodeGeneratorResponse
|
// Determine the output filename
|
||||||
// We'll put all generated code into a single file for now.
|
|
||||||
// If file_to_generate is provided, we'll use the first one as a base for the name.
|
|
||||||
let mut output_filename = "roto_generated.rs".to_string();
|
let mut output_filename = "roto_generated.rs".to_string();
|
||||||
if let Some(first_file) = request.file_to_generate().next() {
|
if let Some(first_file) = request.file_to_generate().next() {
|
||||||
if let Ok((name_bytes, _)) = first_file {
|
if let Ok((name_bytes, _)) = first_file {
|
||||||
@@ -78,6 +81,7 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Construct the response
|
||||||
let mut response_buf = vec![0u8; 1024 * 1024 * 2]; // Allocate 2MB for response
|
let mut response_buf = vec![0u8; 1024 * 1024 * 2]; // Allocate 2MB for response
|
||||||
let mut resp_builder = CodeGeneratorResponse::builder(&mut response_buf);
|
let mut resp_builder = CodeGeneratorResponse::builder(&mut response_buf);
|
||||||
|
|
||||||
@@ -91,7 +95,7 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|||||||
e
|
e
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let final_response = resp_builder
|
let final_response_slice = resp_builder
|
||||||
.add_file(final_file)?
|
.add_file(final_file)?
|
||||||
.finish()
|
.finish()
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
@@ -99,9 +103,32 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
|
|||||||
e
|
e
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
// 5. Write response to stdout
|
// The finish() method returns a reference to the buffer.
|
||||||
io::stdout().write_all(final_response)?;
|
// We convert it to a Vec<u8> to return it from this function.
|
||||||
|
Ok(final_response_slice.to_vec())
|
||||||
info!("Successfully wrote CodeGeneratorResponse to stdout");
|
}
|
||||||
Ok(())
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use std::fs;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_handle_request_with_bin() {
|
||||||
|
// Note: this test assumes request.bin exists in the current working directory
|
||||||
|
// which is usually the project root during 'cargo test'.
|
||||||
|
let request_path = "request.bin";
|
||||||
|
if !std::path::Path::new(request_path).exists() {
|
||||||
|
return; // Skip if file is not available in the environment
|
||||||
|
}
|
||||||
|
|
||||||
|
let data = fs::read(request_path).expect("Failed to read request.bin");
|
||||||
|
let request = CodeGeneratorRequest::new(&data).expect("Failed to parse CodeGeneratorRequest");
|
||||||
|
|
||||||
|
let result = handle_request(&request);
|
||||||
|
|
||||||
|
assert!(result.is_ok(), "handle_request should succeed with request.bin");
|
||||||
|
let response = result.unwrap();
|
||||||
|
assert!(!response.is_empty(), "The generated response should not be empty");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user