diff --git a/src/main.rs b/src/main.rs
index e78fe10..ac25ebf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,4 @@
-use std::io::{stdin, stdout, Read, Write};
+use std::io::{stdin, stdout, Read, Stdout, Write};
 
 use crossterm::{cursor, event, terminal, ExecutableCommand, Result};
 use event::{Event, KeyCode, KeyEvent};
@@ -248,19 +248,24 @@ impl State {
     }
 }
 
-fn main() -> Result<()> {
-    let text = include_str!("main.rs");
-    let mut state = State::from_str(text);
-    let mut stdout = stdout();
-    terminal::enable_raw_mode()?;
-    stdout.execute(terminal::EnterAlternateScreen)?;
-
+fn screen_main(stdout: &mut Stdout, mut state: State) -> Result<()> {
     while !state.quit {
-        state.draw(&mut stdout)?;
+        state.draw(stdout)?;
         let event = event::read()?;
         state.on_event(event);
     }
 
-    stdout.execute(terminal::LeaveAlternateScreen)?;
     Ok(())
 }
+
+fn main() -> Result<()> {
+    let text = include_str!("main.rs");
+    let state = State::from_str(text);
+    let mut stdout = stdout();
+    terminal::enable_raw_mode()?;
+    stdout.execute(terminal::EnterAlternateScreen)?;
+    let result = screen_main(&mut stdout, state);
+    stdout.execute(terminal::LeaveAlternateScreen)?;
+    terminal::disable_raw_mode()?;
+    result
+}