1 files changed, 38 insertions, 3 deletions

diff --git a/src/main.rs b/src/main.rs
index d1d4ee4..03c8115 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@ mod app;
 mod logging;
 mod ui;
 
-use app::App;
+use app::{App, AppMode};
 use miette::{IntoDiagnostic, Result, miette};
 use ratatui::{Terminal, backend::Backend};
 
@@ -77,16 +77,27 @@ fn handle_key_event(app: &mut App, key_event: KeyEvent) -> Result<()> {
     tracing::event!(
         tracing::Level::TRACE,
         event = key_event.code.to_string(),
+        mode = app.get_app_mode().to_string(),
         "Received key event."
     );
+    match app.get_app_mode() {
+        AppMode::Counter => handle_counter_key_event(app, key_event),
+        // no keyevent handling for init, exit modes
+        AppMode::Init => Ok(()),
+        AppMode::Exit => Ok(()),
+    }?;
+    // any new key event, we should clear the existing error
+    app.clear_error()?;
+    Ok(())
+}
+
+fn handle_counter_key_event(app: &mut App, key_event: KeyEvent) -> Result<()> {
     match key_event.code {
         KeyCode::Char('q') => app.exit()?,
         KeyCode::Left => app.decrement_counter()?,
         KeyCode::Right => app.increment_counter()?,
         _ => {}
     };
-    // any new key event, we should clear the existing error
-    app.clear_error()?;
     Ok(())
 }
 
@@ -97,6 +108,30 @@ mod tests {
     use app::OverflowError;
 
     #[test]
+    fn handle_key_event_app_init() {
+        let mut app = App::default();
+        let res = handle_key_event(&mut app, KeyCode::Right.into());
+        assert!(res.is_ok());
+        assert_eq!(*app.get_app_mode(), AppMode::Init);
+
+        // Function returns "Error"; app never init'd
+        assert_eq!(app.display_counter(), "Error");
+    }
+
+    #[test]
+    fn handle_key_event_app_exit() {
+        let mut app = App::default();
+        assert!(app.exit().is_ok());
+
+        let res = handle_key_event(&mut app, KeyCode::Right.into());
+        assert!(res.is_ok());
+        assert_eq!(*app.get_app_mode(), AppMode::Exit);
+
+        // Function returns "Error"; app never init'd
+        assert_eq!(app.display_counter(), "Error");
+    }
+
+    #[test]
     fn handle_key_event_inc_dec() {
         let mut app = App::default();
         assert!(app.init_context().is_ok());