From 187a4712cb858a1fdb06ffc692ffa7398f453552 Mon Sep 17 00:00:00 2001
From: Aayush Goel <81844215+Aayush-Goel-04@users.noreply.github.com>
Date: Wed, 10 May 2023 14:39:16 +0530
Subject: [PATCH] Update test_scripts.py

Here new_rule_path and expected_overlaps will be changed based on the new test rule designed.
Adding tests to check if the code works fine
---
 scripts/detect_duplicate_features.py | 41 ++++++++++++++++------------
 tests/test_scripts.py                | 12 ++++++++
 2 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/scripts/detect_duplicate_features.py b/scripts/detect_duplicate_features.py
index 0ce8ff3a..a72688c0 100644
--- a/scripts/detect_duplicate_features.py
+++ b/scripts/detect_duplicate_features.py
@@ -1,9 +1,13 @@
+import sys
+import logging
 import argparse
 
 import capa.main
 import capa.rules
 import capa.engine as ceng
 
+logger = logging.getLogger("detect_duplicate_features")
+
 
 def get_child_features(feature: ceng.Statement) -> list:
     """
@@ -43,13 +47,15 @@ def get_features(rule_path: str) -> list:
             new_rule = capa.rules.Rule.from_yaml(f.read())
             feature_list = get_child_features(new_rule.statement)
         except Exception as e:
-            raise Warning("Error: " + rule_path + " " + str(type(e)) + " " + str(e))
+            logger.error("Error: New rule " + rule_path + " " + str(type(e)) + " " + str(e))
+            sys.exit(1)
     return feature_list
 
 
 def find_overlapping_rules(new_rule_path, rules_path):
     if not new_rule_path.endswith(".yml"):
-        raise FileNotFoundError("FileNotFoundError ! New rule file name doesn't end with yml")
+        logger.error("FileNotFoundError ! New rule file name doesn't end with .yml")
+        sys.exit(1)
 
     # Loads features of new rule in a list.
     new_rule_features = get_features(new_rule_path)
@@ -84,21 +90,22 @@ def main():
 
     new_rule_path = args.new_rule
     rules_path = args.rules
-    try:
-        result = find_overlapping_rules(new_rule_path, rules_path)
-        print("\nNew rule path : %s" % new_rule_path)
-        print("Number of rules checked : %s " % result["count"])
-        if result["overlapping_rules"]:
-            print("Paths to overlapping rules : ")
-            for r in result["overlapping_rules"]:
-                print("- %s" % r)
-        else:
-            print("Paths to overlapping rules : None")
-        print("Number of rules containing same features : %s" % len(result["overlapping_rules"]))
-        print("\n")
-    except Exception as e:
-        print(e)
+
+    result = find_overlapping_rules(new_rule_path, rules_path)
+
+    print("\nNew rule path : %s" % new_rule_path)
+    print("Number of rules checked : %s " % result["count"])
+    if result["overlapping_rules"]:
+        print("Paths to overlapping rules : ")
+        for r in result["overlapping_rules"]:
+            print("- %s" % r)
+    else:
+        print("Paths to overlapping rules : None")
+    print("Number of rules containing same features : %s" % len(result["overlapping_rules"]))
+    print("\n")
+
+    return len(result["overlapping_rules"])
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())
diff --git a/tests/test_scripts.py b/tests/test_scripts.py
index 3b299a36..eca6a1f4 100644
--- a/tests/test_scripts.py
+++ b/tests/test_scripts.py
@@ -82,3 +82,15 @@ def test_proto_conversion(tmpdir):
     assert p.returncode == 0
 
     assert p.stdout.startswith(b'{\n  "meta": ') or p.stdout.startswith(b'{\r\n  "meta": ')
+
+
+def test_detect_duplicate_features():
+    new_rule_path = "collection/credit-card/parse-credit-card-information.yml"
+    args = [
+        get_rules_path(),
+        os.path.join(get_rules_path(), new_rule_path),
+    ]
+    expected_overlaps = 49
+    script_path = get_script_path("detect_duplicate_features.py")
+    p = run_program(script_path, args)
+    assert p.returncode == expected_overlaps