From 2bd4bf3a716811348342431d1cdd7bc847ec02ff Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Tue, 7 Apr 2026 05:11:30 +0000
Subject: [PATCH] test: add edge case tests for Engine.greedyNextToken and
 decodeStep

This commit adds a test for the `decodeStep` function and significantly expands the existing tests for `Engine.greedyNextToken` in `echo-core-zig/src/inference/engine.zig`. We now explicitly test how `greedyNextToken` handles situations like all negative logits and ties.

Co-authored-by: ulac000000 <132948319+ulac000000@users.noreply.github.com>
---
 echo-core-zig/src/inference/engine.zig | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/echo-core-zig/src/inference/engine.zig b/echo-core-zig/src/inference/engine.zig
index eae5fd0..cd08f1e 100644
--- a/echo-core-zig/src/inference/engine.zig
+++ b/echo-core-zig/src/inference/engine.zig
@@ -1015,4 +1015,27 @@ test "Engine.greedyNextToken correctly identifies token with max logit" {
     // Test early in the slice
     eng.logits[0] = 10.0;
     try std.testing.expectEqual(@as(u32, 0), eng.greedyNextToken());
+
+    // Edge case: all negative logits
+    eng.logits[0] = -5.0;
+    eng.logits[1] = -10.0;
+    eng.logits[2] = -2.0; // max is -2.0 at index 2
+    eng.logits[3] = -3.0;
+    try std.testing.expectEqual(@as(u32, 2), eng.greedyNextToken());
+
+    // Edge case: tie (should return the first one encountered)
+    eng.logits[0] = 1.0;
+    eng.logits[1] = 3.0; // max
+    eng.logits[2] = 3.0; // tie
+    eng.logits[3] = 2.0;
+    try std.testing.expectEqual(@as(u32, 1), eng.greedyNextToken());
+}
+
+test "Engine.decodeStep returns logits of expected length" {
+    const cfg = makeTinyConfig(0, 0);
+    var eng = try Engine.init(cfg, null, std.testing.allocator);
+    defer eng.deinit(std.testing.allocator);
+
+    const logits = try eng.decodeStep(0);
+    try std.testing.expectEqual(@as(usize, cfg.vocab_size), logits.len);
 }