添加单元测试

RRFRRF · Jan 1, 2025 · 09a99e6 · 09a99e6
1 parent d12ff9b
commit 09a99e6
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 6 deletions.
diff --git a/logo_recog.py b/logo_recog.py
@@ -80,8 +80,8 @@ def vis(img_path, pred_boxes):
     # draw rectangle
     for j, box in enumerate(pred_boxes):
         if j == 0:
-            cv2.rectangle(check, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), COLORS['0'], 2)
+            cv2.rectangle(check, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), COLORS[0], 2)
         else:
-            cv2.rectangle(check, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), COLORS['1'], 2)
+            cv2.rectangle(check, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), COLORS[1], 2)
 
     return check
diff --git a/tests/__pycache__/test_logo_matching.cpython-38-pytest-8.3.4.pyc b/tests/__pycache__/test_logo_matching.cpython-38-pytest-8.3.4.pyc
diff --git a/tests/__pycache__/test_logo_recog.cpython-38-pytest-8.3.4.pyc b/tests/__pycache__/test_logo_recog.cpython-38-pytest-8.3.4.pyc
diff --git a/tests/__pycache__/test_phishpedia.cpython-38-pytest-8.3.4.pyc b/tests/__pycache__/test_phishpedia.cpython-38-pytest-8.3.4.pyc
diff --git a/tests/test_logo_matching.py b/tests/test_logo_matching.py
@@ -50,10 +50,12 @@ def test_no_logo_boxes():
 
 # 测试用例 2：有 logo 框但 pred_brand 不返回匹配
 def test_logo_boxes_no_match():
-    with patch("logo_matching.tldextract.extract") as mock_extract, \
+    with patch("builtins.open", mock_open(read_data=pickle.dumps(domain_map_sample))) as mock_file, \
+         patch("logo_matching.tldextract.extract") as mock_extract, \
          patch("logo_matching.pred_brand") as mock_pred_brand, \
          patch("logo_matching.brand_converter") as mock_brand_converter:
 
+        mock_file = mock_file
         # 设置 tldextract.extract 返回值
         mock_extract.return_value = MagicMock(domain="example", suffix="com")
 
@@ -83,10 +85,12 @@ def test_logo_boxes_no_match():
 
 # 测试用例 3：有 logo 框且域名一致
 def test_logo_boxes_domain_consistent():
-    with patch("logo_matching.tldextract.extract") as mock_extract, \
+    with patch("builtins.open", mock_open(read_data=pickle.dumps(domain_map_sample))) as mock_file, \
+         patch("logo_matching.tldextract.extract") as mock_extract, \
          patch("logo_matching.pred_brand") as mock_pred_brand, \
          patch("logo_matching.brand_converter") as mock_brand_converter:
 
+        mock_file = mock_file
         # 设置 tldextract.extract 返回值
         mock_extract.return_value = MagicMock(domain="example", suffix="com")
 
@@ -127,10 +131,12 @@ def brand_converter_side_effect(arg):
 
 # 测试用例 4：有 logo 框且域名不一致
 def test_logo_boxes_domain_inconsistent():
-    with patch("logo_matching.tldextract.extract") as mock_extract, \
+    with patch("builtins.open", mock_open(read_data=pickle.dumps(domain_map_sample))) as mock_file, \
+         patch("logo_matching.tldextract.extract") as mock_extract, \
          patch("logo_matching.pred_brand") as mock_pred_brand, \
          patch("logo_matching.brand_converter") as mock_brand_converter:
 
+        mock_file = mock_file
         # 设置 tldextract.extract 返回值
         mock_extract.return_value = MagicMock(domain="example", suffix="com")
 
@@ -160,10 +166,12 @@ def test_logo_boxes_domain_inconsistent():
 
 # 测试用例 5：超过 topk 的 logo 框
 def test_logo_boxes_exceed_topk():
-    with patch("logo_matching.tldextract.extract") as mock_extract, \
+    with patch("builtins.open", mock_open(read_data=pickle.dumps(domain_map_sample))) as mock_file, \
+         patch("logo_matching.tldextract.extract") as mock_extract, \
          patch("logo_matching.pred_brand") as mock_pred_brand, \
          patch("logo_matching.brand_converter") as mock_brand_converter:
 
+        mock_file = mock_file
         # 设置 tldextract.extract 返回值
         mock_extract.return_value = MagicMock(domain="example", suffix="com")
 

diff --git a/tests/test_logo_recog.py b/tests/test_logo_recog.py
@@ -33,9 +33,12 @@ def test_pred_rcnn_image_with_alpha_channel():
             mock_instances.pred_classes.__eq__.return_value = MagicMock()
             mock_instances.pred_boxes.__getitem__.return_value = MagicMock()
 
+            result = pred_rcnn('path/to/image_with_alpha.png', mock_predictor)
+
             mock_imread.assert_called_once_with('path/to/image_with_alpha.png')
             mock_cvtColor.assert_called_once()
             mock_predictor.assert_called_once()
+            assert result
             # 根据 mock 的返回值，断言结果
             # 由于我们没有具体的返回值设置，这里仅断言返回值被调用
             # 在实际测试中，您应根据具体的 mock 返回值进行断言