From 70a53bf12b8710da7a2a2372419a92d052af7972 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 21 Nov 2022 13:31:17 -0500
Subject: [PATCH] Add unit tests for complexity check (#859)

---
 src/mccabe/checks.rs | 232 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 232 insertions(+)

diff --git a/src/mccabe/checks.rs b/src/mccabe/checks.rs
index 2d20fa98bb..204ffc7fd3 100644
--- a/src/mccabe/checks.rs
+++ b/src/mccabe/checks.rs
@@ -70,3 +70,235 @@ pub fn function_is_too_complex(
         None
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use rustpython_parser::parser;
+
+    use crate::mccabe::checks::get_complexity_number;
+
+    #[test]
+    fn trivial() -> Result<()> {
+        let source = r#"
+def trivial():
+    pass
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 1);
+        Ok(())
+    }
+
+    #[test]
+    fn expr_as_statement() -> Result<()> {
+        let source = r#"
+def expr_as_statement():
+    0xF00D
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 1);
+        Ok(())
+    }
+
+    #[test]
+    fn sequential() -> Result<()> {
+        let source = r#"
+def sequential(n):
+    k = n + 4
+    s = k + n
+    return s
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 1);
+        Ok(())
+    }
+
+    #[test]
+    fn if_elif_else_dead_path() -> Result<()> {
+        let source = r#"
+def if_elif_else_dead_path(n):
+    if n > 3:
+        return "bigger than three"
+    elif n > 4:
+        return "is never executed"
+    else:
+        return "smaller than or equal to three"
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 3);
+        Ok(())
+    }
+
+    #[test]
+    fn nested_ifs() -> Result<()> {
+        let source = r#"
+def nested_ifs():
+    if n > 3:
+        if n > 4:
+            return "bigger than four"
+        else:
+            return "bigger than three"
+    else:
+        return "smaller than or equal to three"
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 3);
+        Ok(())
+    }
+
+    #[test]
+    fn for_loop() -> Result<()> {
+        let source = r#"
+def for_loop():
+    for i in range(10):
+        print(i)
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn for_else() -> Result<()> {
+        let source = r#"
+def for_else(mylist):
+    for i in mylist:
+        print(i)
+    else:
+        print(None)
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn recursive() -> Result<()> {
+        let source = r#"
+def recursive(n):
+    if n > 4:
+        return f(n - 1)
+    else:
+        return n
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn nested_functions() -> Result<()> {
+        let source = r#"
+def nested_functions():
+    def a():
+        def b():
+            pass
+
+        b()
+
+    a()
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 3);
+        Ok(())
+    }
+
+    #[test]
+    fn try_else() -> Result<()> {
+        let source = r#"
+def try_else():
+    try:
+        print(1)
+    except TypeA:
+        print(2)
+    except TypeB:
+        print(3)
+    else:
+        print(4)
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 4);
+        Ok(())
+    }
+
+    #[test]
+    fn nested_try_finally() -> Result<()> {
+        let source = r#"
+def nested_try_finally():
+    try:
+        try:
+            print(1)
+        finally:
+            print(2)
+    finally:
+        print(3)
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 3);
+        Ok(())
+    }
+
+    #[test]
+    fn foobar() -> Result<()> {
+        let source = r#"
+async def foobar(a, b, c):
+    await whatever(a, b, c)
+    if await b:
+        pass
+    async with c:
+        pass
+    async for x in a:
+        pass
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 3);
+        Ok(())
+    }
+
+    #[test]
+    fn annotated_assign() -> Result<()> {
+        let source = r#"
+def annotated_assign():
+    x: Any = None
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 1);
+        Ok(())
+    }
+
+    #[test]
+    fn class() -> Result<()> {
+        let source = r#"
+class Class:
+    def handle(self, *args, **options):
+        if args:
+            return
+
+        class ServiceProvider:
+            def a(self):
+                pass
+
+            def b(self, data):
+                if not args:
+                    pass
+
+        class Logger:
+            def c(*args, **kwargs):
+                pass
+
+            def error(self, message):
+                pass
+
+            def info(self, message):
+                pass
+
+            def exception(self):
+                pass
+
+        return ServiceProvider(Logger())
+"#;
+        let stmts = parser::parse_program(source, "<filename>")?;
+        assert_eq!(get_complexity_number(&stmts), 9);
+        Ok(())
+    }
+}