diff --git a/crates/tabby/src/serve/completions/prompt.rs b/crates/tabby/src/serve/completions/prompt.rs
index ba9fe94..6c827c5 100644
--- a/crates/tabby/src/serve/completions/prompt.rs
+++ b/crates/tabby/src/serve/completions/prompt.rs
@@ -200,3 +200,156 @@ lazy_static! {
     static ref LANGUAGE_LINE_COMMENT_CHAR: HashMap<&'static str, &'static str> =
         HashMap::from([("python", "#"), ("rust", "//"),]);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn create_prompt_builder(with_template: bool) -> PromptBuilder {
+        let prompt_template = if with_template {
+            // Init prompt builder with codellama prompt template
+            Some("<PRE> {prefix} <SUF>{suffix} <MID>".into())
+        } else {
+            None
+        };
+
+        // Init prompt builder with prompt rewrite disabled.
+        PromptBuilder::new(prompt_template, false)
+    }
+
+    #[test]
+    fn test_prompt_template() {
+        let pb = create_prompt_builder(true);
+
+        // Rewrite disabled, so the language doesn't matter.
+        let language = "python";
+
+        // Test w/ prefix, w/ suffix.
+        {
+            let segments = Segments {
+                prefix: "this is some prefix".into(),
+                suffix: Some("this is some suffix".into()),
+            };
+            assert_eq!(
+                pb.build(language, segments),
+                "<PRE> this is some prefix <SUF>this is some suffix <MID>"
+            );
+        }
+
+        // Test w/ prefix, w/o suffix.
+        {
+            let segments = Segments {
+                prefix: "this is some prefix".into(),
+                suffix: None,
+            };
+            assert_eq!(
+                pb.build(language, segments),
+                "<PRE> this is some prefix <SUF>\n <MID>"
+            );
+        }
+
+        // Test w/ prefix, w/ empty suffix.
+        {
+            let segments = Segments {
+                prefix: "this is some prefix".into(),
+                suffix: Some("".into()),
+            };
+            assert_eq!(
+                pb.build(language, segments),
+                "<PRE> this is some prefix <SUF> <MID>"
+            );
+        }
+
+        // Test w/ empty prefix, w/ suffix.
+        {
+            let segments = Segments {
+                prefix: "".into(),
+                suffix: Some("this is some suffix".into()),
+            };
+            assert_eq!(
+                pb.build(language, segments),
+                "<PRE>  <SUF>this is some suffix <MID>"
+            );
+        }
+
+        // Test w/ empty prefix, w/o suffix.
+        {
+            let segments = Segments {
+                prefix: "".into(),
+                suffix: None,
+            };
+            assert_eq!(pb.build(language, segments), "<PRE>  <SUF>\n <MID>");
+        }
+
+        // Test w/ emtpy prefix, w/ empty suffix.
+        {
+            let segments = Segments {
+                prefix: "".into(),
+                suffix: Some("".into()),
+            };
+            assert_eq!(pb.build(language, segments), "<PRE>  <SUF> <MID>");
+        }
+    }
+
+    #[test]
+    fn test_no_prompt_template() {
+        let pb = create_prompt_builder(false);
+
+        // Rewrite disabled, so the language doesn't matter.
+        let language = "python";
+
+        // Test w/ prefix, w/ suffix.
+        {
+            let segments = Segments {
+                prefix: "this is some prefix".into(),
+                suffix: Some("this is some suffix".into()),
+            };
+            assert_eq!(pb.build(language, segments), "this is some prefix");
+        }
+
+        // Test w/ prefix, w/o suffix.
+        {
+            let segments = Segments {
+                prefix: "this is some prefix".into(),
+                suffix: None,
+            };
+            assert_eq!(pb.build(language, segments), "this is some prefix");
+        }
+
+        // Test w/ prefix, w/ empty suffix.
+        {
+            let segments = Segments {
+                prefix: "this is some prefix".into(),
+                suffix: Some("".into()),
+            };
+            assert_eq!(pb.build(language, segments), "this is some prefix");
+        }
+
+        // Test w/ empty prefix, w/ suffix.
+        {
+            let segments = Segments {
+                prefix: "".into(),
+                suffix: Some("this is some suffix".into()),
+            };
+            assert_eq!(pb.build(language, segments), "");
+        }
+
+        // Test w/ empty prefix, w/o suffix.
+        {
+            let segments = Segments {
+                prefix: "".into(),
+                suffix: None,
+            };
+            assert_eq!(pb.build(language, segments), "");
+        }
+
+        // Test w/ empty prefix, w/ empty suffix.
+        {
+            let segments = Segments {
+                prefix: "".into(),
+                suffix: Some("".into()),
+            };
+            assert_eq!(pb.build(language, segments), "");
+        }
+    }
+}