2 files changed, 70 insertions, 7 deletions

diff --git a/src/app/mod.rs b/src/app/mod.rs
index d51640f..7bb9c72 100644
--- a/src/app/mod.rs
+++ b/src/app/mod.rs
@@ -5,11 +5,33 @@ use miette::{Diagnostic, ErrReport, Result};
 use store::{Context, Store};
 use thiserror::Error;
 
+#[derive(Debug, Default, PartialEq)]
+pub enum AppMode {
+    #[default]
+    Init,
+    Counter,
+    Exit,
+}
+
+impl std::fmt::Display for AppMode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}",
+            match self {
+                AppMode::Init => "INIT",
+                AppMode::Counter => "COUNTER",
+                AppMode::Exit => "EXIT",
+            }
+        )
+    }
+}
+
 #[derive(Default)]
 pub struct App {
     context: Context,
     error: Option<ErrReport>,
-    exit: bool,
+    mode: AppMode,
 }
 
 #[derive(Diagnostic, Debug, Error, PartialEq)]
@@ -26,16 +48,22 @@ impl App {
         let store = self.context.unwrap_store();
 
         store.init()?;
+
+        self.mode = AppMode::Counter;
         Ok(self)
     }
 
+    pub fn get_app_mode(&self) -> &AppMode {
+        &self.mode
+    }
+
     pub fn get_exit_status(&self) -> bool {
-        self.exit
+        self.mode == AppMode::Exit
     }
 
     pub fn exit(&mut self) -> Result<()> {
         tracing::event!(tracing::Level::TRACE, exit = true, "Changing exit status.");
-        self.exit = true;
+        self.mode = AppMode::Exit;
         Ok(())
     }
 
@@ -183,7 +211,7 @@ mod tests {
         let res = app.exit();
 
         assert!(res.is_ok());
-        assert!(app.exit);
+        assert_eq!(app.mode, AppMode::Exit);
     }
 
     #[test]
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());