diff --git a/src/main.rs b/src/main.rs
index 18e03650..739b1983 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,11 @@
 #![no_std]
-#![no_main]
+#![cfg_attr(not(test), no_main)]
 
 use core::panic::PanicInfo;
 
 mod vga_buffer;
 
+#[cfg(not(test))]
 #[no_mangle]
 pub extern "C" fn _start() -> ! {
     println!("Hello World{}", "!");
@@ -13,6 +14,7 @@ pub extern "C" fn _start() -> ! {
 }
 
 /// This function is called on panic.
+#[cfg(not(test))]
 #[panic_handler]
 fn panic(info: &PanicInfo) -> ! {
     println!("{}", info);