From 6ded7b53f7f47ab77beae2718b48070af28d62f0 Mon Sep 17 00:00:00 2001
From: James Smith <james@foxglove.dev>
Date: Tue, 7 Jan 2025 13:49:57 +1100
Subject: [PATCH] add errors and tests

---
 rust/src/lib.rs   |  4 +--
 rust/src/write.rs | 80 +++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 261c67b5e..6d5457b90 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -140,9 +140,9 @@ pub enum McapError {
     ChunkBufferTooLarge(u64),
     #[error("length exceeds usize max: `{0}`")]
     TooLong(u64),
-    #[error("cannot write more than 65335 channels to one MCAP")]
+    #[error("cannot write more than 65536 channels to one MCAP")]
     TooManyChannels,
-    #[error("cannot write more than 65334 schemas to one MCAP")]
+    #[error("cannot write more than 65535 schemas to one MCAP")]
     TooManySchemas,
 }
 
diff --git a/rust/src/write.rs b/rust/src/write.rs
index 48bd04774..4a9a4c9dd 100644
--- a/rust/src/write.rs
+++ b/rust/src/write.rs
@@ -306,9 +306,6 @@ impl<W: Write + Seek> Writer<W> {
             },
             schema.id,
         );
-        if schema.id >= self.next_schema_id {
-            self.next_schema_id = schema.id + 1;
-        }
         if self.options.use_chunks {
             self.chunkin_time()?.write_schema(schema)
         } else {
@@ -386,9 +383,6 @@ impl<W: Write + Seek> Writer<W> {
             },
             channel.id,
         );
-        if channel.id >= self.next_channel_id {
-            self.next_channel_id = channel.id + 1;
-        }
         if self.options.use_chunks {
             self.chunkin_time()?.write_channel(channel)
         } else {
@@ -1306,3 +1300,77 @@ impl<W: Write + Seek> AttachmentWriter<W> {
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::u16;
+
+    use super::*;
+    #[test]
+    fn writes_all_channel_ids() {
+        let file = std::io::Cursor::new(Vec::new());
+        let mut writer = Writer::new(file).expect("failed to construct writer");
+        let custom_channel = std::sync::Arc::new(crate::Channel {
+            id: u16::MAX,
+            topic: "chat".into(),
+            message_encoding: "json".into(),
+            metadata: BTreeMap::new(),
+            schema: None,
+        });
+        writer
+            .write(&crate::Message {
+                channel: custom_channel.clone(),
+                sequence: 0,
+                log_time: 0,
+                publish_time: 0,
+                data: Cow::Owned(Vec::new()),
+            })
+            .expect("could not write initial channel");
+        for i in 0..65535u16 {
+            let id = writer
+                .add_channel(0, &format!("{i}"), "json", &BTreeMap::new())
+                .expect("could not add channel");
+            assert_eq!(i, id);
+        }
+        let Err(too_many) = writer.add_channel(0, "last", "json", &BTreeMap::new()) else {
+            panic!("should not be able to add another channel");
+        };
+        assert!(matches!(too_many, McapError::TooManyChannels));
+    }
+    #[test]
+    fn writes_all_schema_ids() {
+        let file = std::io::Cursor::new(Vec::new());
+        let mut writer = Writer::new(file).expect("failed to construct writer");
+        let custom_channel = std::sync::Arc::new(crate::Channel {
+            id: 0,
+            topic: "chat".into(),
+            message_encoding: "json".into(),
+            metadata: BTreeMap::new(),
+            schema: Some(std::sync::Arc::new(crate::Schema {
+                id: u16::MAX,
+                name: "int".into(),
+                encoding: "jsonschema".into(),
+                data: Cow::Owned(Vec::new()),
+            })),
+        });
+        writer
+            .write(&crate::Message {
+                channel: custom_channel.clone(),
+                sequence: 0,
+                log_time: 0,
+                publish_time: 0,
+                data: Cow::Owned(Vec::new()),
+            })
+            .expect("could not write initial channel");
+        for i in 0..65534u16 {
+            let id = writer
+                .add_schema(&format!("{i}"), "jsonschema", &[])
+                .expect("could not add schema");
+            assert_eq!(id, i + 1);
+        }
+        let Err(too_many) = writer.add_schema("last", "jsonschema", &[]) else {
+            panic!("should not be able to add another channel");
+        };
+        assert!(matches!(too_many, McapError::TooManySchemas));
+    }
+}