From fbcae960469f63de103ff8088678d328cb3898f5 Mon Sep 17 00:00:00 2001
From: Itamar Gafni <itamar@intezer.com>
Date: Tue, 13 Aug 2024 16:03:34 +0300
Subject: [PATCH] Add unit tests for write_bytecode_file

---
 pytest/test_std.py                            | 58 +++++++++++++++++++
 pytest/testdata/03_annotations-xasm-3.7.right |  1 +
 2 files changed, 59 insertions(+)

diff --git a/pytest/test_std.py b/pytest/test_std.py
index 7080feeb..a34eb0a8 100644
--- a/pytest/test_std.py
+++ b/pytest/test_std.py
@@ -1,6 +1,10 @@
 # std
+import os
+import shutil
 import sys
+import tempfile
 from contextlib import closing
+from datetime import datetime
 
 # compat
 import six
@@ -11,6 +15,10 @@
 # local
 import xdis.std as dis
 from xdis import IS_GRAAL, IS_PYPY, PYTHON3, PYTHON_VERSION_TRIPLE
+from xdis import Code3
+from xdis import list2bytecode
+from xdis import opcodes
+from xdis import write_bytecode_file
 
 if PYTHON_VERSION_TRIPLE >= (3, 2):
     if pytest.__version__ >= "3.2.0":
@@ -175,6 +183,56 @@ def test_findlabels():
         actual_len = len(actual)
         assert actual_len > 0
 
+    def _create_python3_code_object():
+        consts = (None, 2)
+        varnames = ("a",)
+        instructions = [
+            ("LOAD_CONST", 2),
+            ("STORE_FAST", "a"),
+            ("LOAD_FAST", "a"),
+            ("RETURN_VALUE",),
+        ]
+
+        text = "def fact():\n\ta = 8\n\ta = 0\n"
+
+        code = list2bytecode(instructions, opcodes.opcode_34, varnames, consts)
+        fn_code = compile(text, "<string>", "exec")
+        return Code3(
+            fn_code.co_argcount,
+            fn_code.co_kwonlyargcount,  # Add this in Python3
+            fn_code.co_nlocals,
+            fn_code.co_stacksize,
+            fn_code.co_flags,
+            # These are changed
+            code,
+            consts,
+            fn_code.co_names,
+            varnames,
+            fn_code.co_filename,
+            fn_code.co_name,
+            fn_code.co_firstlineno,
+            fn_code.co_lnotab,  # In general, You should adjust this
+            fn_code.co_freevars,
+            fn_code.co_cellvars,
+        )
+
+    def test_write_bytecode_file():
+        temp_dir = tempfile.mkdtemp()
+        target_path = os.path.join(temp_dir, "test1.pyc")
+        code_object = _create_python3_code_object()
+        write_bytecode_file(target_path, code_object, 3394, 10)
+        shutil.rmtree(temp_dir)
+
+    def test_write_bytecode_bad_timestamp_type():
+        temp_dir = tempfile.mkdtemp()
+        target_path = os.path.join(temp_dir, "test2.pyc")
+        code_object = _create_python3_code_object()
+        with pytest.raises(TypeError):
+            write_bytecode_file(
+                target_path, code_object, 3394, datetime.now().timestamp()
+            )
+        shutil.rmtree(temp_dir)
+
 
 if __name__ == "__main__":
     test_findlabels()
diff --git a/pytest/testdata/03_annotations-xasm-3.7.right b/pytest/testdata/03_annotations-xasm-3.7.right
index 83dfb7f8..fb8c193d 100644
--- a/pytest/testdata/03_annotations-xasm-3.7.right
+++ b/pytest/testdata/03_annotations-xasm-3.7.right
@@ -24,3 +24,4 @@
             POP_TOP
             LOAD_CONST           2 (None)
             RETURN_VALUE
+