open-metadata--openmetadata/ingestion/tests/unit/readers/test_json_reader.py

#  Copyright 2025 Collate
#  Licensed under the Collate Community License, Version 1.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
#  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.

"""
Tests for JSONDataFrameReader
"""

import gzip
import json
import tempfile
import unittest
import zipfile
from unittest.mock import Mock, patch

from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
    AzureConfig,
)
from metadata.generated.schema.entity.services.connections.database.datalake.gcsConfig import (
    GCSConfig,
)
from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import (
    S3Config,
)
from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
    LocalConfig,
)
from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
from metadata.generated.schema.security.credentials.azureCredentials import (
    AzureCredentials,
)
from metadata.readers.dataframe.json import JSONDataFrameReader


class TestJSONReader(unittest.TestCase):
    def test_json_lines_local(self):
        json_lines = '{"id": 1, "name": "Alice"}\n{"id": 2, "name": "Bob"}\n'

        with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as tmp:
            tmp.write(json_lines)
            tmp_path = tmp.name

        try:
            config = LocalConfig()
            reader = JSONDataFrameReader(config, None)

            result = reader._read(key=tmp_path, bucket_name="")

            self.assertIsNotNone(result.dataframes)
            dataframes = result.dataframes()
            self.assertIsNotNone(dataframes)
            chunks = list(dataframes)
            self.assertTrue(len(chunks) > 0)
            total_rows = sum(len(chunk) for chunk in chunks)
            self.assertEqual(total_rows, 2)
        finally:
            import os

            os.unlink(tmp_path)  # noqa: PTH108

    def test_json_array_local(self):
        json_array = [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]

        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
            json.dump(json_array, tmp)
            tmp_path = tmp.name

        try:
            config = LocalConfig()
            reader = JSONDataFrameReader(config, None)

            result = reader._read(key=tmp_path, bucket_name="")

            self.assertIsNotNone(result.dataframes)
            dataframes = result.dataframes()
            self.assertIsNotNone(dataframes)
            chunks = list(dataframes)
            total_rows = sum(len(chunk) for chunk in chunks)
            self.assertEqual(total_rows, 2)
        finally:
            import os

            os.unlink(tmp_path)  # noqa: PTH108

    def test_json_object_local(self):
        json_obj = {"id": 1, "name": "Alice"}

        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
            json.dump(json_obj, tmp)
            tmp_path = tmp.name

        try:
            config = LocalConfig()
            reader = JSONDataFrameReader(config, None)

            result = reader._read(key=tmp_path, bucket_name="")

            self.assertIsNotNone(result.dataframes)
            dataframes = result.dataframes()
            self.assertIsNotNone(dataframes)
            chunks = list(dataframes)
            self.assertEqual(len(chunks), 1)
            self.assertEqual(chunks[0].shape[0], 1)
        finally:
            import os

            os.unlink(tmp_path)  # noqa: PTH108

    def test_json_gzip_compression(self):
        json_data = [{"id": 1, "name": "Test"}]

        with tempfile.NamedTemporaryFile(suffix=".json.gz", delete=False) as tmp:
            with gzip.open(tmp.name, "wt") as gz:
                json.dump(json_data, gz)
            tmp_path = tmp.name

        try:
            config = LocalConfig()
            reader = JSONDataFrameReader(config, None)

            result = reader._read(key=tmp_path, bucket_name="")

            self.assertIsNotNone(result.dataframes)
            dataframes = result.dataframes()
            self.assertIsNotNone(dataframes)
            chunks = list(dataframes)
            self.assertTrue(len(chunks) > 0)
        finally:
            import os

            os.unlink(tmp_path)  # noqa: PTH108

    def test_json_zip_compression(self):
        json_data = [{"id": 1, "name": "Test"}]

        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
            tmp_path = tmp.name

        try:
            with zipfile.ZipFile(tmp_path, "w") as zf:
                zf.writestr("data.json", json.dumps(json_data))

            config = LocalConfig()
            reader = JSONDataFrameReader(config, None)

            result = reader._read(key=tmp_path, bucket_name="")

            self.assertIsNotNone(result.dataframes)
            dataframes = result.dataframes()
            self.assertIsNotNone(dataframes)
            chunks = list(dataframes)
            self.assertTrue(len(chunks) > 0)
        finally:
            import os

            os.unlink(tmp_path)  # noqa: PTH108

    def test_is_json_lines_detection(self):
        import io

        jsonl_content = '{"id": 1}\n{"id": 2}\n'
        jsonl_file = io.BytesIO(jsonl_content.encode())

        is_jsonl = JSONDataFrameReader._is_json_lines(jsonl_file)
        self.assertTrue(is_jsonl)

    def test_stream_json_lines(self):
        import io

        jsonl_content = '{"id": 1}\n{"id": 2}\n{"id": 3}\n'
        jsonl_file = io.BytesIO(jsonl_content.encode())

        chunks = list(JSONDataFrameReader._stream_json_lines(jsonl_file, batch_size=2))

        self.assertEqual(len(chunks), 2)
        self.assertEqual(len(chunks[0]), 2)
        self.assertEqual(len(chunks[1]), 1)

    @patch("gcsfs.GCSFileSystem")
    def test_gcs_json_reading(self, mock_gcsfs):
        import io

        mock_gcs = Mock()
        mock_gcsfs.return_value = mock_gcs

        json_data = b'{"id": 1, "name": "Test"}'
        mock_file = io.BytesIO(json_data)
        mock_gcs.open.return_value.__enter__ = Mock(return_value=mock_file)
        mock_gcs.open.return_value.__exit__ = Mock(return_value=False)
        mock_gcs.info.return_value = {"size": len(json_data)}

        config = GCSConfig()
        reader = JSONDataFrameReader(config, None)

        result = reader._read(key="test.json", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)
        dataframes = result.dataframes()
        self.assertIsNotNone(dataframes)

    @patch("adlfs.AzureBlobFileSystem")
    @patch("metadata.readers.dataframe.json.return_azure_storage_options")
    def test_azure_json_reading(self, mock_storage_opts, mock_adlfs):
        import io

        mock_storage_opts.return_value = {"connection_string": "test"}
        mock_fs = Mock()
        mock_adlfs.return_value = mock_fs

        json_data = b'{"id": 1, "name": "Test"}'
        mock_file = io.BytesIO(json_data)
        mock_fs.open.return_value.__enter__ = Mock(return_value=mock_file)
        mock_fs.open.return_value.__exit__ = Mock(return_value=False)
        mock_fs.info.return_value = {"size": len(json_data)}

        config = AzureConfig(securityConfig=AzureCredentials(accountName="test", clientId="test", tenantId="test"))
        reader = JSONDataFrameReader(config, None)

        result = reader._read(key="test.json", bucket_name="test-container")

        self.assertIsNotNone(result.dataframes)
        dataframes = result.dataframes()
        self.assertIsNotNone(dataframes)

    def test_s3_json_reading(self):
        import io

        json_data = b'{"id": 1, "name": "Test"}'

        mock_client = Mock()
        mock_body = io.BytesIO(json_data)
        mock_client.get_object.return_value = {"Body": mock_body}
        mock_client.head_object.return_value = {"ContentLength": len(json_data)}

        config = S3Config(
            securityConfig=AWSCredentials(awsAccessKeyId="test", awsSecretAccessKey="test", awsRegion="us-east-1")
        )
        reader = JSONDataFrameReader(config, mock_client)

        result = reader._read(key="test.json", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)
        dataframes = result.dataframes()
        self.assertIsNotNone(dataframes)

    def test_empty_json_lines(self):
        import io

        jsonl_content = '{"id": 1}\n\n{"id": 2}\n'
        jsonl_file = io.BytesIO(jsonl_content.encode())

        chunks = list(JSONDataFrameReader._stream_json_lines(jsonl_file))

        total_rows = sum(len(chunk) for chunk in chunks)
        self.assertEqual(total_rows, 2)


if __name__ == "__main__":
    unittest.main()