From 02fdf419693ec3b4d553e027d8a103b5a24a99c7 Mon Sep 17 00:00:00 2001
From: Willi Ballenthin <willi.ballenthin@gmail.com>
Date: Wed, 22 Mar 2023 10:47:45 +0100
Subject: [PATCH] tests: add tests demonstrating result document round tripping

---
 tests/fixtures.py             | 13 +++++------
 tests/test_result_document.py | 41 +++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/tests/fixtures.py b/tests/fixtures.py
index 7f26c888..b8c572b4 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -1115,30 +1115,29 @@ def get_result_doc(path):
 
 @pytest.fixture
 def pma0101_rd():
-    # TODO move to rd subdir
-    return get_result_doc(os.path.join(CD, "data", "Practical Malware Analysis Lab 01-01.dll_.json"))
+    return get_result_doc(os.path.join(CD, "data", "rd", "Practical Malware Analysis Lab 01-01.dll_.json"))
 
 
 @pytest.fixture
 def dotnet_1c444e_rd():
-    return get_result_doc(os.path.join(CD, "data", "dotnet", "1c444ebeba24dcba8628b7dfe5fec7c6.exe_.json"))
+    return get_result_doc(os.path.join(CD, "data", "rd", "1c444ebeba24dcba8628b7dfe5fec7c6.exe_.json"))
 
 
 @pytest.fixture
 def a3f3bbc_rd():
-    return get_result_doc(os.path.join(CD, "data", "3f3bbcf8fd90bdcdcdc5494314ed4225.exe_.json"))
+    return get_result_doc(os.path.join(CD, "data", "rd", "3f3bbcf8fd90bdcdcdc5494314ed4225.exe_.json"))
 
 
 @pytest.fixture
 def al_khaserx86_rd():
-    return get_result_doc(os.path.join(CD, "data", "al-khaser_x86.exe_.json"))
+    return get_result_doc(os.path.join(CD, "data", "rd", "al-khaser_x86.exe_.json"))
 
 
 @pytest.fixture
 def al_khaserx64_rd():
-    return get_result_doc(os.path.join(CD, "data", "al-khaser_x64.exe_.json"))
+    return get_result_doc(os.path.join(CD, "data", "rd", "al-khaser_x64.exe_.json"))
 
 
 @pytest.fixture
 def a076114_rd():
-    return get_result_doc(os.path.join(CD, "data", "0761142efbda6c4b1e801223de723578.dll_.json"))
+    return get_result_doc(os.path.join(CD, "data", "rd", "0761142efbda6c4b1e801223de723578.dll_.json"))
diff --git a/tests/test_result_document.py b/tests/test_result_document.py
index 8172c601..976fd0da 100644
--- a/tests/test_result_document.py
+++ b/tests/test_result_document.py
@@ -5,12 +5,16 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 #  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and limitations under the License.
+import copy
+import pytest
 
 import capa
 import capa.engine as ceng
 import capa.render.result_document as rdoc
 import capa.features.freeze.features as frzf
 
+from fixtures import *
+
 
 def test_optional_node_from_capa():
     node = rdoc.node_from_capa(
@@ -227,3 +231,40 @@ def test_basic_block_node_from_capa():
     node = rdoc.node_from_capa(capa.features.basicblock.BasicBlock(""))
     assert isinstance(node, rdoc.FeatureNode)
     assert isinstance(node.feature, frzf.BasicBlockFeature)
+
+
+def assert_round_trip(rd: rdoc.ResultDocument):
+    one = rd
+
+    doc = one.json(exclude_none=True)
+    two = rdoc.ResultDocument.parse_raw(doc)
+
+    # show the round trip works
+    # first by comparing the objects directly,
+    # which works thanks to pydantic model equality.
+    assert one == two
+    # second by showing their json representations are the same.
+    assert one.json(exclude_none=True) == two.json(exclude_none=True)
+
+    # now show that two different versions are not equal.
+    three = copy.deepcopy(two)
+    three.meta.__dict__.update({"version": "0.0.0"})
+    assert one.meta.version != three.meta.version
+    assert one != three
+    assert one.json(exclude_none=True) != three.json(exclude_none=True)
+
+
+@pytest.mark.parametrize(
+    "rd_file",
+    [
+        pytest.param("a3f3bbc_rd"),
+        pytest.param("al_khaserx86_rd"),
+        pytest.param("al_khaserx64_rd"),
+        pytest.param("a076114_rd"),
+        pytest.param("pma0101_rd"),
+        pytest.param("dotnet_1c444e_rd"),
+    ],
+)
+def test_round_trip(request, rd_file):
+    rd: rdoc.ResultDocument = request.getfixturevalue(rd_file)
+    assert_round_trip(rd)
\ No newline at end of file