Re-write json2numpy tests

Signed-off-by: Jeremy Johnson <jeremy.johnson@arm.com>
Change-Id: I17227adfca203f353e51ae9297f7e01f8d7edbe9

1 files changed, 64 insertions, 82 deletions

diff --git a/verif/tests/test_json2numpy.py b/verif/tests/test_json2numpy.py
index b8ea48d..b01ebe9 100644
--- a/verif/tests/test_json2numpy.py
+++ b/verif/tests/test_json2numpy.py
@@ -1,7 +1,6 @@
 """Tests for json2numpy.py."""
-# Copyright (c) 2021-2022, ARM Limited.
+# Copyright (c) 2021-2023, ARM Limited.
 # SPDX-License-Identifier: Apache-2.0
-import json
 import os
 
 import numpy as np
@@ -9,6 +8,18 @@ import pytest
 from json2numpy.json2numpy import main
 
 
+DTYPE_RANGES = {
+    np.int8.__name__: [-128, 128],
+    np.uint8.__name__: [0, 256],
+    np.int16.__name__: [-32768, 32768],
+    np.uint16.__name__: [0, 65536],
+    np.int32.__name__: [-(1 << 31), (1 << 31)],
+    np.uint32.__name__: [0, (1 << 32)],
+    np.int64.__name__: [-(1 << 63), (1 << 63)],
+    np.uint64.__name__: [0, (1 << 64)],
+}
+
+
 @pytest.mark.parametrize(
     "npy_filename,json_filename,data_type",
     [
@@ -26,10 +37,11 @@ from json2numpy.json2numpy import main
         ("multiple_num.npy", "multiple_num.json", np.uint16),
         ("single_num.npy", "single_num.json", np.uint32),
         ("multiple_num.npy", "multiple_num.json", np.uint32),
-        ("single_num.npy", "single_num.json", np.uint64),
-        ("multiple_num.npy", "multiple_num.json", np.uint64),
-        #        ("single_num.npy", "single_num.json", np.float16),
-        #        ("multiple_num.npy", "multiple_num.json", np.float16),
+        # Not implemented due to json.dump issue
+        # ("single_num.npy", "single_num.json", np.uint64),
+        # ("multiple_num.npy", "multiple_num.json", np.uint64),
+        ("single_num.npy", "single_num.json", np.float16),
+        ("multiple_num.npy", "multiple_num.json", np.float16),
         ("single_num.npy", "single_num.json", np.float32),
         ("multiple_num.npy", "multiple_num.json", np.float32),
         ("single_num.npy", "single_num.json", np.float64),
@@ -38,13 +50,34 @@ from json2numpy.json2numpy import main
         ("multiple_num.npy", "multiple_num.json", bool),
     ],
 )
-def test_json2numpy_npy_file(npy_filename, json_filename, data_type):
+def test_json2numpy_there_and_back(npy_filename, json_filename, data_type):
     """Test conversion to JSON."""
     # Generate numpy data.
     if "single" in npy_filename:
-        npy_data = np.ndarray(shape=(1, 1), dtype=data_type)
+        shape = (1,)
     elif "multiple" in npy_filename:
-        npy_data = np.ndarray(shape=(2, 3), dtype=data_type)
+        shape = (4, 6, 5)
+
+    rng = np.random.default_rng()
+    nan_location = None
+    if data_type in [np.float16, np.float32, np.float64]:
+        gen_type = np.float32 if data_type == np.float16 else data_type
+        generated_npy_data = rng.standard_normal(size=shape, dtype=gen_type).astype(
+            data_type
+        )
+        if len(shape) > 1:
+            # Set some NANs and INFs
+            nan_location = (1, 2, 3)
+            generated_npy_data[nan_location] = np.nan
+            generated_npy_data[(3, 2, 1)] = np.inf
+            generated_npy_data[(0, 5, 2)] = -np.inf
+    elif data_type == bool:
+        generated_npy_data = rng.choice([True, False], size=shape).astype(bool)
+    else:
+        range = DTYPE_RANGES[data_type.__name__]
+        generated_npy_data = rng.integers(
+            low=range[0], high=range[1], size=shape, dtype=data_type
+        )
 
     # Get filepaths
     npy_file = os.path.join(os.path.dirname(__file__), npy_filename)
@@ -52,87 +85,36 @@ def test_json2numpy_npy_file(npy_filename, json_filename, data_type):
 
     # Save npy data to file and reload it.
     with open(npy_file, "wb") as f:
-        np.save(f, npy_data)
+        np.save(f, generated_npy_data)
     npy_data = np.load(npy_file)
 
+    # Test json2numpy - converts npy file to json
     args = [npy_file]
-    """Converts npy file to json"""
     assert main(args) == 0
 
-    json_data = json.load(open(json_file))
-    assert np.dtype(json_data["type"]) == npy_data.dtype
-    assert np.array(json_data["data"]).shape == npy_data.shape
-    assert (np.array(json_data["data"], dtype=data_type) == npy_data).all()
-
-    # Remove files created
-    if os.path.exists(npy_file):
-        os.remove(npy_file)
-    if os.path.exists(json_file):
-        os.remove(json_file)
-
-
-@pytest.mark.parametrize(
-    "npy_filename,json_filename,data_type",
-    [
-        ("single_num.npy", "single_num.json", np.int8),
-        ("multiple_num.npy", "multiple_num.json", np.int8),
-        ("single_num.npy", "single_num.json", np.int16),
-        ("multiple_num.npy", "multiple_num.json", np.int16),
-        ("single_num.npy", "single_num.json", np.int32),
-        ("multiple_num.npy", "multiple_num.json", np.int32),
-        ("single_num.npy", "single_num.json", np.int64),
-        ("multiple_num.npy", "multiple_num.json", np.int64),
-        ("single_num.npy", "single_num.json", np.uint8),
-        ("multiple_num.npy", "multiple_num.json", np.uint8),
-        ("single_num.npy", "single_num.json", np.uint16),
-        ("multiple_num.npy", "multiple_num.json", np.uint16),
-        ("single_num.npy", "single_num.json", np.uint32),
-        ("multiple_num.npy", "multiple_num.json", np.uint32),
-        ("single_num.npy", "single_num.json", np.uint64),
-        ("multiple_num.npy", "multiple_num.json", np.uint64),
-        #        ("single_num.npy", "single_num.json", np.float16),
-        #        ("multiple_num.npy", "multiple_num.json", np.float16),
-        ("single_num.npy", "single_num.json", np.float32),
-        ("multiple_num.npy", "multiple_num.json", np.float32),
-        ("single_num.npy", "single_num.json", np.float64),
-        ("multiple_num.npy", "multiple_num.json", np.float64),
-        ("single_num.npy", "single_num.json", bool),
-        ("multiple_num.npy", "multiple_num.json", bool),
-    ],
-)
-def test_json2numpy_json_file(npy_filename, json_filename, data_type):
-    """Test conversion to binary."""
-    # Generate json data.
-    if "single" in npy_filename:
-        npy_data = np.ndarray(shape=(1, 1), dtype=data_type)
-    elif "multiple" in npy_filename:
-        npy_data = np.ndarray(shape=(2, 3), dtype=data_type)
-
-    # Generate json dictionary
-    list_data = npy_data.tolist()
-    json_data_type = str(npy_data.dtype)
-
-    json_data = {}
-    json_data["type"] = json_data_type
-    json_data["data"] = list_data
-
-    # Get filepaths
-    npy_file = os.path.join(os.path.dirname(__file__), npy_filename)
-    json_file = os.path.join(os.path.dirname(__file__), json_filename)
-
-    # Save json data to file and reload it.
-    with open(json_file, "w") as f:
-        json.dump(json_data, f)
-    json_data = json.load(open(json_file))
-
+    # Remove the numpy file and convert json back to npy
+    os.remove(npy_file)
+    assert not os.path.exists(npy_file)
     args = [json_file]
-    """Converts json file to npy"""
     assert main(args) == 0
 
-    npy_data = np.load(npy_file)
-    assert np.dtype(json_data["type"]) == npy_data.dtype
-    assert np.array(json_data["data"]).shape == npy_data.shape
-    assert (np.array(json_data["data"], dtype=data_type) == npy_data).all()
+    converted_npy_data = np.load(npy_file)
+
+    # Check that the original data equals the npy->json->npy data
+    assert converted_npy_data.dtype == npy_data.dtype
+    assert converted_npy_data.shape == npy_data.shape
+    equals = np.equal(converted_npy_data, npy_data)
+    if nan_location is not None:
+        # NaNs do not usaually equal - so check and set
+        if np.isnan(converted_npy_data[nan_location]) and np.isnan(
+            npy_data[nan_location]
+        ):
+            equals[nan_location] = True
+    if not np.all(equals):
+        print("JSONed:  ", converted_npy_data)
+        print("Original:", npy_data)
+        print("Equals:  ", equals)
+    assert np.all(equals)
 
     # Remove files created
     if os.path.exists(npy_file):