From 6b8bce4f4247fc397dfd5ca29e693eb55639de92 Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Fri, 14 Aug 2020 12:34:00 -0600
Subject: [PATCH] tests: fixtures: factor out resolution of scope/sample

---
 tests/fixtures.py          | 30 ++++++++++++++++++----------
 tests/test_ida_features.py | 41 ++++++++++++++++++++++----------------
 2 files changed, 43 insertions(+), 28 deletions(-)

diff --git a/tests/fixtures.py b/tests/fixtures.py
index 385415d7..d7c69431 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -104,9 +104,13 @@ def get_data_path_by_name(name):
         raise ValueError("unexpected sample fixture")
 
 
+def resolve_sample(sample):
+    return get_data_path_by_name(request.param)
+
+
 @pytest.fixture
 def sample(request):
-    return get_data_path_by_name(request.param)
+    return resolve_sample(request.param)
 
 
 def get_function(extractor, fva):
@@ -123,18 +127,17 @@ def get_basic_block(extractor, f, va):
     raise ValueError("basic block not found")
 
 
-@pytest.fixture
-def scope(request):
-    if request.param == "file":
+def resolve_scope(scope):
+    if scope == "file":
 
         def inner(extractor):
             return extract_file_features(extractor)
 
-        inner.__name__ = request.param
+        inner.__name__ = scope
         return inner
-    elif "bb=" in request.param:
+    elif "bb=" in scope:
         # like `function=0x401000,bb=0x40100A`
-        fspec, _, bbspec = request.param.partition(",")
+        fspec, _, bbspec = scope.partition(",")
         fva = int(fspec.partition("=")[2], 0x10)
         bbva = int(bbspec.partition("=")[2], 0x10)
 
@@ -143,22 +146,27 @@ def scope(request):
             bb = get_basic_block(extractor, f, bbva)
             return extract_basic_block_features(extractor, f, bb)
 
-        inner.__name__ = request.param
+        inner.__name__ = scope
         return inner
-    elif request.param.startswith("function"):
+    elif scope.startswith("function"):
         # like `function=0x401000`
-        va = int(request.param.partition("=")[2], 0x10)
+        va = int(scope.partition("=")[2], 0x10)
 
         def inner(extractor):
             f = get_function(extractor, va)
             return extract_function_features(extractor, f)
 
-        inner.__name__ = request.param
+        inner.__name__ = scope
         return inner
     else:
         raise ValueError("unexpected scope fixture")
 
 
+@pytest.fixture
+def scope(request):
+    return resolve_scope(request.param)
+
+
 def make_test_id(values):
     return "-".join(map(str, values))
 
diff --git a/tests/test_ida_features.py b/tests/test_ida_features.py
index 2156375e..8161d1cc 100644
--- a/tests/test_ida_features.py
+++ b/tests/test_ida_features.py
@@ -1,17 +1,18 @@
 # run this script from within IDA with ./tests/data/mimikatz.exe open
+import sys
 import logging
+import os.path
 import binascii
-import traceback
-import collections
 
 import pytest
-from fixtures import *
 
-import capa.features
-import capa.features.file
-import capa.features.insn
-import capa.features.basicblock
-from capa.features import ARCH_X32, ARCH_X64
+try:
+    sys.path.append(os.path.dirname(__file__))
+    print(sys.path)
+    from fixtures import *
+finally:
+    sys.path.pop()
+
 
 logger = logging.getLogger("test_ida_features")
 
@@ -43,11 +44,9 @@ def get_ida_extractor(_path):
 @pytest.mark.skip(reason="IDA Pro tests must be run within IDA")
 def test_ida_features():
     for (sample, scope, feature, expected) in FEATURE_PRESENCE_TESTS:
-        # resolve sample
-        # resolve scope
-        pass
-
         id = make_test_id((sample, scope, feature, expected))
+        scope = resolve_scope(scope)
+        sample = resolve_sample(sample)
 
         try:
             do_test_feature_presence(get_ida_extractor, sample, scope, feature, expected)
@@ -59,11 +58,19 @@ def test_ida_features():
 
 
 @pytest.mark.skip(reason="IDA Pro tests must be run within IDA")
-@parametrize(
-    "sample,scope,feature,expected", FEATURE_COUNT_TESTS, indirect=["sample", "scope"],
-)
-def test_ida_feature_counts(sample, scope, feature, expected):
-    do_test_feature_count(get_ida_extractor, sample, scope, feature, expected)
+def test_ida_feature_counts():
+    for (sample, scope, feature, expected) in FEATURE_COUNT_TESTS:
+        id = make_test_id((sample, scope, feature, expected))
+        scope = resolve_scope(scope)
+        sample = resolve_sample(sample)
+
+        try:
+            do_test_feature_count(get_ida_extractor, sample, scope, feature, expected)
+        except AssertionError as e:
+            print("FAIL %s" % (id))
+            traceback.print_exc()
+        else:
+            print("OK   %s" % (id))
 
 
 if __name__ == "__main__":