avdb-py/tests/test_db.py

import unittest
import os
import numpy as np
import csv
import uuid
from unittest.mock import MagicMock, patch
import psycopg2
from psycopg2 import extras

from avdb.db import DBOps
from avdb.analysis import convert_mmap_to_csv

class TestDBOps(unittest.TestCase):
    def setUp(self):
        # Mock the database connection and cursor
        self.mock_conn = MagicMock()
        self.mock_cursor = MagicMock()
        self.mock_conn.get_cursor.return_value = self.mock_cursor
        self.mock_conn.get_connection.return_value = self.mock_conn

        # Initialize DBOps with the mocked connection
        self.db_ops = DBOps(self.mock_conn)

    def tearDown(self):
        # Clean up any temporary files created during tests
        if os.path.exists("test_histogram.mmap"):
            os.remove("test_histogram.mmap")
        if os.path.exists("test_histogram.csv"):
            os.remove("test_histogram.csv")
        if os.path.exists("test_tile.mmap"):
            os.remove("test_tile.mmap")

    def test_convert_mmap_to_csv(self):
        # Create a temporary .mmap file for testing
        mmap_file_path = "test_histogram.mmap"
        csv_file_path = "test_histogram.csv"

        # Create dummy data for the .mmap file
        dummy_data = np.random.rand(81, 48).astype('float32')
        with open(mmap_file_path, 'wb') as f:
            f.write(dummy_data.tobytes())

        # Call the conversion function
        convert_mmap_to_csv(mmap_file_path, csv_file_path)

        # Verify the CSV file was created and contains the correct data
        with open(csv_file_path, 'r') as csv_file:
            reader = csv.reader(csv_file)
            csv_data = list(reader)

        # Convert the CSV data back to a numpy array for comparison
        csv_data_array = np.array(csv_data, dtype='float32')

        # Assert that the data matches
        np.testing.assert_array_equal(csv_data_array, dummy_data)

    def test_insert_tile_data(self):
        # Create dummy tile data
        dummy_rows = [
            (
                1,  # img_id
                3,  # level
                0,  # tilex
                0,  # tiley
                100,  # x
                100,  # y
                50,  # x0
                50,  # y0
                10,  # w
                10,  # h
                str(uuid.uuid4()),  # hash
                [1, 2, 3]  # tile_data
            ),
            (
                1,  # img_id
                3,  # level
                1,  # tilex
                1,  # tiley
                200,  # x
                200,  # y
                100,  # x0
                100,  # y0
                10,  # w
                10,  # h
                str(uuid.uuid4()),  # hash
                [4, 5, 6]  # tile_data
            )
        ]

        # Mock the execute_batch method
        with patch('avdb.db.psycopg2.extras.execute_batch') as mock_execute_batch:
            # Call the insert_tile_data method
            self.db_ops.insert_tile_data(dummy_rows)

            # Verify that execute_batch was called with the correct query and data
            mock_execute_batch.assert_called_once_with(
                self.mock_cursor,
                """
                INSERT INTO pyramiddata_level3 (
                    img_id, level, tilex, tiley, x, y, x0, y0, w, h, hash, tile_data
                )
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """,
                dummy_rows
            )
        
    def test_process_and_insert_tile_data_batched(self):
        # Create a temporary .mmap file for testing
        mmap_file_path = "test_tile.mmap"
        dummy_data = np.random.randint(0, 255, size=(100, 3), dtype=np.uint16)
        with open(mmap_file_path, 'wb') as f:
            f.write(dummy_data.tobytes())

        # Mock the load_memmapped_tile_data method to return the dummy data
        self.db_ops.load_memmapped_tile_data = MagicMock(return_value=dummy_data)

        # Mock the execute_batch method
        with patch('avdb.db.psycopg2.extras.execute_batch') as mock_execute_batch:
            # Call the process_and_insert_tile_data_batched method
            self.db_ops.process_and_insert_tile_data_batched(
                img_id=1,
                img_dims=(1000, 1000),
                level=3,
                tile_file=mmap_file_path
            )

            # Verify that execute_batch was called
            mock_execute_batch.assert_called()

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