142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
import unittest
|
|
import json
|
|
import struct
|
|
import os
|
|
from unittest.mock import patch, mock_open
|
|
|
|
from amqp.service.stream_parser import StreamParser
|
|
|
|
|
|
class TestStreamParser(unittest.TestCase):
|
|
def setUp(self):
|
|
"""
|
|
Set up the StreamParser instance for testing.
|
|
"""
|
|
self.parser = StreamParser()
|
|
|
|
def tearDown(self):
|
|
"""
|
|
Clean up after each test.
|
|
"""
|
|
if hasattr(self.parser, 'file_writer') and self.parser.file_writer:
|
|
self.parser.file_writer.close()
|
|
if os.path.exists("test_output_file"):
|
|
os.remove("test_output_file")
|
|
|
|
def test_process_chunk_with_metadata_and_data(self):
|
|
"""
|
|
Test processing a chunk that contains metadata and data.
|
|
"""
|
|
# Create test data
|
|
metadata = {"filename": "test_output_file", "description": "Test file"}
|
|
metadata_json = json.dumps(metadata).encode('utf-8')
|
|
metadata_length = struct.pack('>I', len(metadata_json))
|
|
data = b"some binary data"
|
|
|
|
# Combine metadata length, metadata, and data into a single chunk
|
|
chunk = metadata_length + metadata_json + data
|
|
|
|
# Process the chunk
|
|
self.parser.process_chunk(chunk)
|
|
self.parser.close()
|
|
|
|
# Verify metadata
|
|
self.assertEqual(self.parser.metadata, metadata)
|
|
self.assertEqual(self.parser.metadata_length, len(metadata_json))
|
|
|
|
# Verify file content
|
|
with open("test_output_file", "rb") as f:
|
|
file_content = f.read()
|
|
self.assertEqual(file_content, data)
|
|
|
|
def test_process_chunk_in_parts(self):
|
|
"""
|
|
Test processing chunks in multiple parts.
|
|
"""
|
|
# Create test data
|
|
metadata = {"filename": "test_output_file", "description": "Test file"}
|
|
metadata_json = json.dumps(metadata).encode('utf-8')
|
|
metadata_length = struct.pack('>I', len(metadata_json))
|
|
data = b"some binary data"
|
|
|
|
# Split the chunk into parts
|
|
chunk1 = metadata_length # First 4 bytes (metadata length)
|
|
chunk2 = metadata_json # Metadata JSON
|
|
chunk3 = data # Remaining data
|
|
|
|
# Process chunks one by one
|
|
self.parser.process_chunk(chunk1)
|
|
self.parser.process_chunk(chunk2)
|
|
self.parser.process_chunk(chunk3)
|
|
self.parser.close()
|
|
|
|
# Verify metadata
|
|
self.assertEqual(self.parser.metadata, metadata)
|
|
self.assertEqual(self.parser.metadata_length, len(metadata_json))
|
|
|
|
# Verify file content
|
|
with open("test_output_file", "rb") as f:
|
|
file_content = f.read()
|
|
self.assertEqual(file_content, data)
|
|
|
|
def test_process_chunk_without_metadata(self):
|
|
"""
|
|
Test processing a chunk without metadata.
|
|
"""
|
|
# Create test data (only metadata length and metadata, no additional data)
|
|
metadata = {"filename": "distinct_test_output_file", "description": "Test file"}
|
|
metadata_json = json.dumps(metadata).encode('utf-8')
|
|
metadata_length = struct.pack('>I', len(metadata_json))
|
|
|
|
# Combine metadata length and metadata into a single chunk
|
|
chunk = metadata_length + metadata_json
|
|
|
|
# Process the chunk
|
|
self.parser.process_chunk(chunk)
|
|
self.parser.close()
|
|
|
|
# Verify metadata
|
|
self.assertEqual(self.parser.metadata, metadata)
|
|
self.assertEqual(self.parser.metadata_length, len(metadata_json))
|
|
|
|
# Verify no data was written to the file
|
|
self.assertFalse(os.path.exists("distinct_test_output_file"))
|
|
|
|
def test_process_chunk_with_invalid_metadata(self):
|
|
"""
|
|
Test processing a chunk with invalid metadata.
|
|
"""
|
|
# Create invalid metadata (not JSON)
|
|
metadata_length = struct.pack('>I', 10)
|
|
invalid_metadata = b"invalid_json"
|
|
|
|
# Combine metadata length and invalid metadata into a single chunk
|
|
chunk = metadata_length + invalid_metadata
|
|
|
|
# Process the chunk and expect an exception
|
|
with self.assertRaises(json.JSONDecodeError):
|
|
self.parser.process_chunk(chunk)
|
|
|
|
|
|
|
|
@patch("builtins.open", new_callable=mock_open)
|
|
def test_file_writing(self, mock_file):
|
|
"""
|
|
Test file writing functionality using a mock file.
|
|
"""
|
|
# Create test data
|
|
metadata = {"filename": "test_output_file", "description": "Test file"}
|
|
metadata_json = json.dumps(metadata).encode('utf-8')
|
|
metadata_length = struct.pack('>I', len(metadata_json))
|
|
data = b"some binary data"
|
|
|
|
# Combine metadata length, metadata, and data into a single chunk
|
|
chunk = metadata_length + metadata_json + data
|
|
|
|
# Process the chunk
|
|
self.parser.process_chunk(chunk)
|
|
|
|
# Verify the file was opened and written to
|
|
mock_file.assert_called_once_with("test_output_file", "wb")
|
|
mock_file().write.assert_called_once_with(data)
|