Add some tests

src/bin/generator.rs

@@ -1,6 +1,6 @@
 use clap::Parser;
 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::path::PathBuf;
 

src/bin/protoc-gen-roto.rs

@@ -29,46 +29,49 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
 
     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
-    // Since generate_rust_code expects a FileDescriptorSet, we wrap the proto_files in one.
-    // A FileDescriptorSet is a message with a repeated field 'file' at tag 1.
-    let mut set_buf = vec![0u8; 1024 * 1024]; // Allocate 1MB for descriptors
-    {
-        let mut builder = ProtoBuilder::new(&mut set_buf);
-        for file_res in request.proto_file() {
-            let (file_data, _) = file_res.map_err(|e| {
-                error!("Failed to iterate proto_file: {:?}", e);
-                e
-            })?;
-            builder.write_bytes(1, file_data).map_err(|e| {
-                error!("Failed to write proto_file to set: {:?}", e);
-                e
-            })?;
-        }
-        // We don't need to call finish() here as we just need the buffer to be populated.
+    let mut set_buf = Vec::new();
+    for file_res in request.proto_file() {
+        let (file_data, _) = file_res.map_err(|e| {
+            error!("Failed to iterate proto_file: {:?}", e);
+            e
+        })?;
+
+        // Tag 1, Length-delimited: (1 << 3) | 2 = 10
+        set_buf.push(10);
+
+        // Write length as varint
+        let len = file_data.len() as u64;
+        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);
+            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)?;
 
-    // 3. Generate the Rust code
+    // Generate the Rust code
     info!("Generating Rust code from descriptor set...");
     let generated_code = generate_rust_code(&set);
 
-    // 4. Construct CodeGeneratorResponse
-    // 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.
+    // Determine the output filename
     let mut output_filename = "roto_generated.rs".to_string();
     if let Some(first_file) = request.file_to_generate().next() {
         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 resp_builder = CodeGeneratorResponse::builder(&mut response_buf);
 
@@ -91,7 +95,7 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
             e
         })?;
 
-    let final_response = resp_builder
+    let final_response_slice = resp_builder
         .add_file(final_file)?
         .finish()
         .map_err(|e| {
@@ -99,9 +103,32 @@ fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
             e
         })?;
 
-    // 5. Write response to stdout
-    io::stdout().write_all(final_response)?;
-
-    info!("Successfully wrote CodeGeneratorResponse to stdout");
-    Ok(())
+    // The finish() method returns a reference to the buffer.
+    // We convert it to a Vec<u8> to return it from this function.
+    Ok(final_response_slice.to_vec())
+}
+
+#[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");
+    }
 }