Skip to content

Commit

Permalink
update file names in scanned and skipped
Browse files Browse the repository at this point in the history
  • Loading branch information
@swashko
swashko committed Feb 12, 2024
1 parent df3e8d6 commit 4e0ecd7
Showing 1 changed file with 64 additions and 74 deletions.
138 changes: 64 additions & 74 deletions tests/test_modelscan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -455,34 +455,30 @@ def test_scan_zip(zip_file_path: Any) -> None:

ms = ModelScan()
results = ms.scan(f"{zip_file_path}/test.zip")
assert results["summary"]["scanned"]["scanned_files"] == [
f"{zip_file_path}/test.zip:data.pkl"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"test.zip:data.pkl"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert ms.issues.all_issues == expected


def test_scan_pytorch(pytorch_file_path: Any) -> None:
ms = ModelScan()
results = ms.scan(Path(f"{pytorch_file_path}/bad_pytorch.pt"))
assert results["summary"]["scanned"]["scanned_files"] == [
f"{pytorch_file_path}/bad_pytorch.pt"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"bad_pytorch.pt"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert ms.issues.all_issues == []
assert [error.scan_name for error in ms.errors] == ["pytorch"] # type: ignore[attr-defined]

results = ms.scan(Path(f"{pytorch_file_path}/safe_zip_pytorch.pt"))
assert results["summary"]["scanned"]["scanned_files"] == [
f"{pytorch_file_path}/safe_zip_pytorch.pt:safe_zip_pytorch/data.pkl"
f"safe_zip_pytorch.pt:safe_zip_pytorch/data.pkl"
]
assert results["summary"]["skipped"]["skipped_files"] == []
assert ms.issues.all_issues == []
assert results["errors"] == []

results = ms.scan(Path(f"{pytorch_file_path}/safe_old_format_pytorch.pt"))
assert results["summary"]["scanned"]["scanned_files"] == [
f"{pytorch_file_path}/safe_old_format_pytorch.pt"
f"safe_old_format_pytorch.pt"
]
assert results["summary"]["skipped"]["skipped_files"] == []
assert ms.issues.all_issues == []
Expand All @@ -503,7 +499,7 @@ def test_scan_pytorch(pytorch_file_path: Any) -> None:
]
results = ms.scan(unsafe_zip_path)
assert results["summary"]["scanned"]["scanned_files"] == [
f"{pytorch_file_path}/unsafe_zip_pytorch.pt:unsafe_zip_pytorch/data.pkl"
f"unsafe_zip_pytorch.pt:unsafe_zip_pytorch/data.pkl"
]
assert results["summary"]["skipped"]["skipped_files"] == []
assert ms.issues.all_issues == expected
Expand All @@ -514,9 +510,7 @@ def test_scan_numpy(numpy_file_path: Any) -> None:
ms = ModelScan()
results = ms.scan(f"{numpy_file_path}/safe_numpy.npy")
assert ms.issues.all_issues == []
assert results["summary"]["scanned"]["scanned_files"] == [
f"{numpy_file_path}/safe_numpy.npy"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"safe_numpy.npy"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert results["errors"] == []

Expand All @@ -535,9 +529,7 @@ def test_scan_numpy(numpy_file_path: Any) -> None:

results = ms.scan(f"{numpy_file_path}/unsafe_numpy.npy")
compare_results(ms.issues.all_issues, expected)
assert results["summary"]["scanned"]["scanned_files"] == [
f"{numpy_file_path}/unsafe_numpy.npy"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"unsafe_numpy.npy"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert results["errors"] == []

Expand All @@ -546,18 +538,14 @@ def test_scan_file_path(file_path: Any) -> None:
benign_pickle = ModelScan()
results = benign_pickle.scan(Path(f"{file_path}/data/benign0_v3.pkl"))
assert benign_pickle.issues.all_issues == []
assert results["summary"]["scanned"]["scanned_files"] == [
f"{file_path}/data/benign0_v3.pkl"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"benign0_v3.pkl"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert results["errors"] == []

benign_dill = ModelScan()
results = benign_dill.scan(Path(f"{file_path}/data/benign0_v3.dill"))
assert benign_dill.issues.all_issues == []
assert results["summary"]["scanned"]["scanned_files"] == [
f"{file_path}/data/benign0_v3.dill"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"benign0_v3.dill"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert results["errors"] == []

Expand Down Expand Up @@ -606,9 +594,7 @@ def test_scan_file_path(file_path: Any) -> None:
}
results = malicious0.scan(Path(f"{file_path}/data/malicious0.pkl"))
compare_results(malicious0.issues.all_issues, expected_malicious0)
assert results["summary"]["scanned"]["scanned_files"] == [
f"{file_path}/data/malicious0.pkl"
]
assert results["summary"]["scanned"]["scanned_files"] == [f"malicious0.pkl"]
assert results["summary"]["skipped"]["skipped_files"] == []
assert results["errors"] == []

Expand Down Expand Up @@ -1204,34 +1190,34 @@ def test_scan_directory_path(file_path: str) -> None:
results = ms.scan(p)
compare_results(ms.issues.all_issues, expected)
assert set(results["summary"]["scanned"]["scanned_files"]) == {
f"{p}/malicious1.zip:data.pkl",
f"{p}/malicious0.pkl",
f"{p}/malicious3.pkl",
f"{p}/malicious6.pkl",
f"{p}/malicious7.pkl",
f"{p}/malicious8.pkl",
f"{p}/malicious9.pkl",
f"{p}/malicious10.pkl",
f"{p}/malicious11.pkl",
f"{p}/malicious12.pkl",
f"{p}/malicious13.pkl",
f"{p}/malicious1_v0.dill",
f"{p}/malicious1_v3.dill",
f"{p}/malicious1_v4.dill",
f"{p}/malicious4.pickle",
f"{p}/malicious5.pickle",
f"{p}/malicious1_v0.pkl",
f"{p}/malicious1_v3.pkl",
f"{p}/malicious1_v4.pkl",
f"{p}/malicious2_v0.pkl",
f"{p}/malicious2_v3.pkl",
f"{p}/malicious2_v4.pkl",
f"{p}/benign0_v0.pkl",
f"{p}/benign0_v3.pkl",
f"{p}/benign0_v4.pkl",
f"{p}/benign0_v0.dill",
f"{p}/benign0_v3.dill",
f"{p}/benign0_v4.dill",
f"malicious1.zip:data.pkl",
f"malicious0.pkl",
f"malicious3.pkl",
f"malicious6.pkl",
f"malicious7.pkl",
f"malicious8.pkl",
f"malicious9.pkl",
f"malicious10.pkl",
f"malicious11.pkl",
f"malicious12.pkl",
f"malicious13.pkl",
f"malicious1_v0.dill",
f"malicious1_v3.dill",
f"malicious1_v4.dill",
f"malicious4.pickle",
f"malicious5.pickle",
f"malicious1_v0.pkl",
f"malicious1_v3.pkl",
f"malicious1_v4.pkl",
f"malicious2_v0.pkl",
f"malicious2_v3.pkl",
f"malicious2_v4.pkl",
f"benign0_v0.pkl",
f"benign0_v3.pkl",
f"benign0_v4.pkl",
f"benign0_v0.dill",
f"benign0_v3.dill",
f"benign0_v4.dill",
}
assert results["summary"]["skipped"]["skipped_files"] == []
assert results["errors"] == []
Expand Down Expand Up @@ -1259,16 +1245,18 @@ def test_scan_keras(keras_file_path: Any, file_extension: str) -> None:
assert ms.issues.all_issues == []
if file_extension == ".pb":
assert set(results["summary"]["scanned"]["scanned_files"]) == {
f"{safe_filename}/fingerprint.pb",
f"{safe_filename}/keras_metadata.pb",
f"{safe_filename}/saved_model.pb",
f"fingerprint.pb",
f"keras_metadata.pb",
f"saved_model.pb",
}
assert set(results["summary"]["skipped"]["skipped_files"]) == {
f"{safe_filename}/variables/variables.data-00000-of-00001",
f"{safe_filename}/variables/variables.index",
f"variables/variables.data-00000-of-00001",
f"variables/variables.index",
}
else:
assert results["summary"]["scanned"]["scanned_files"] == [safe_filename]
assert results["summary"]["scanned"]["scanned_files"] == [
f"safe{file_extension}"
]
assert results["summary"]["skipped"]["skipped_files"] == []

assert results["errors"] == []
Expand Down Expand Up @@ -1356,16 +1344,18 @@ def test_scan_keras(keras_file_path: Any, file_extension: str) -> None:

if file_extension == ".pb":
assert set(results["summary"]["scanned"]["scanned_files"]) == {
f"{unsafe_filename}/fingerprint.pb",
f"{unsafe_filename}/keras_metadata.pb",
f"{unsafe_filename}/saved_model.pb",
f"fingerprint.pb",
f"keras_metadata.pb",
f"saved_model.pb",
}
assert set(results["summary"]["skipped"]["skipped_files"]) == {
f"{unsafe_filename}/variables/variables.data-00000-of-00001",
f"{unsafe_filename}/variables/variables.index",
f"variables/variables.data-00000-of-00001",
f"variables/variables.index",
}
else:
assert results["summary"]["scanned"]["scanned_files"] == [unsafe_filename]
assert results["summary"]["scanned"]["scanned_files"] == [
f"unsafe{file_extension}"
]
assert results["summary"]["skipped"]["skipped_files"] == []


Expand All @@ -1378,13 +1368,13 @@ def test_scan_tensorflow(tensorflow_file_path: Any) -> None:
results = ms.scan(Path(f"{safe_tensorflow_model_dir}"))
assert ms.issues.all_issues == []
assert set(results["summary"]["scanned"]["scanned_files"]) == {
f"{safe_tensorflow_model_dir}/fingerprint.pb",
f"{safe_tensorflow_model_dir}/keras_metadata.pb",
f"{safe_tensorflow_model_dir}/saved_model.pb",
f"fingerprint.pb",
f"keras_metadata.pb",
f"saved_model.pb",
}
assert set(results["summary"]["skipped"]["skipped_files"]) == {
f"{safe_tensorflow_model_dir}/variables/variables.data-00000-of-00001",
f"{safe_tensorflow_model_dir}/variables/variables.index",
f"variables/variables.data-00000-of-00001",
f"variables/variables.index",
}
assert results["errors"] == []

Expand Down Expand Up @@ -1415,13 +1405,13 @@ def test_scan_tensorflow(tensorflow_file_path: Any) -> None:

assert ms.issues.all_issues == expected
assert set(results["summary"]["scanned"]["scanned_files"]) == {
f"{unsafe_tensorflow_model_dir}/fingerprint.pb",
f"{unsafe_tensorflow_model_dir}/keras_metadata.pb",
f"{unsafe_tensorflow_model_dir}/saved_model.pb",
f"fingerprint.pb",
f"keras_metadata.pb",
f"saved_model.pb",
}
assert set(results["summary"]["skipped"]["skipped_files"]) == {
f"{unsafe_tensorflow_model_dir}/variables/variables.data-00000-of-00001",
f"{unsafe_tensorflow_model_dir}/variables/variables.index",
f"variables/variables.data-00000-of-00001",
f"variables/variables.index",
}
assert results["errors"] == []

Expand Down

0 comments on commit 4e0ecd7

Please sign in to comment.