From 7b0718399cadf7ba5c74cb19ce28ed369c63c015 Mon Sep 17 00:00:00 2001
From: Dennis Schwerdel <schwerdel@googlemail.com>
Date: Sun, 11 Oct 2020 22:47:55 +0200
Subject: [PATCH] More tests

---
 src/crypto/init.rs | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/crypto/init.rs b/src/crypto/init.rs
index a503dbe..bd7c932 100644
--- a/src/crypto/init.rs
+++ b/src/crypto/init.rs
@@ -449,6 +449,7 @@ impl<P: Payload> InitState<P> {
             self.repeat_last_message(out);
             Ok(())
         } else {
+            self.next_stage = CLOSING;
             Err(Error::CryptoInit("Initialization timeout"))
         }
     }
@@ -777,7 +778,20 @@ mod tests {
         }
     }
 
-    // TODO Test: timeout after 5 retries
+    #[test]
+    fn timeout() {
+        let (mut sender, _receiver) = create_pair();
+        let mut out = MsgBuffer::new(8);
+        sender.send_ping(&mut out);
+        assert_eq!(sender.stage(), STAGE_PONG);
+        for _ in 0..5 {
+            out.clear();
+            sender.every_second(&mut out).unwrap();
+        }
+        out.clear();
+        assert!(sender.every_second(&mut out).is_err());
+        assert_eq!(sender.stage(), CLOSING);
+    }
 
     // TODO Test: duplicated message or replay attacks