diff --git a/src/code_generator.rs b/src/code_generator.rs
new file mode 100644
index 0000000..1d6c8b5
--- /dev/null
+++ b/src/code_generator.rs
@@ -0,0 +1,42 @@
+use rand::RngCore;
+use sha2::{Sha256, Digest};
+
+pub struct CodeGenerator {
+    counter: u64,
+    salt: [u8; 32],
+}
+
+impl CodeGenerator {
+    pub fn generate(&mut self) -> String {
+        let count = self.counter;
+        self.counter += 1;
+
+        let mut hasher = Sha256::new();
+        hasher.update(self.salt);
+        hasher.update(count.to_be_bytes());
+
+        format!("{:x}", hasher.finalize())[..6].to_string()
+    }
+}
+
+impl Default for CodeGenerator {
+    fn default() -> Self {
+        let mut salt = [0; 32];
+        rand::thread_rng().fill_bytes(&mut salt);
+        
+        CodeGenerator {
+            counter: 0,
+            salt,
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    #[test]
+    fn test_generate() {
+        let code = CodeGenerator::default().generate();
+        assert_eq!(code.len(), 6);
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 6245837..2054433 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,12 +6,11 @@ use std::collections::HashMap;
 mod message;
 use message::Message;
 
-use rand::RngCore;
+mod code_generator;
+use code_generator::CodeGenerator;
 
 use tungstenite::protocol::Message as WsMessage;
 
-use sha2::{Sha256, Digest};
-
 fn main() {
     let code_generator = Arc::new(Mutex::new(CodeGenerator::default()));
     let server = TcpListener::bind("127.0.0.1:8080").unwrap();
@@ -63,33 +62,3 @@ struct Room {
 struct Player {
     name: String,
 }
-
-struct CodeGenerator {
-    counter: u64,
-    salt: [u8; 32],
-}
-
-impl CodeGenerator {
-    fn generate(&mut self) -> String {
-        let count = self.counter;
-        self.counter += 1;
-
-        let mut hasher = Sha256::new();
-        hasher.update(self.salt);
-        hasher.update(count.to_be_bytes());
-
-        format!("{:x}", hasher.finalize())[..6].to_string()
-    }
-}
-
-impl Default for CodeGenerator {
-    fn default() -> Self {
-        let mut salt = [0; 32];
-        rand::thread_rng().fill_bytes(&mut salt);
-        
-        CodeGenerator {
-            counter: 0,
-            salt,
-        }
-    }
-}
diff --git a/src/message.rs b/src/message.rs
index 7a0494c..87570e4 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -45,7 +45,6 @@ mod test {
 	    command: "COMMAND",
 	    args: vec!["arg1", "arg2"],
 	}, msg);
-
 	Ok(())
     }
 }