test_work_anniv.py

"""
This script is used to test the Qxf2's work anniversary image that is generated by lambda
for each employee on their work anniversary.
It verifies that the year of anniversary displayed in the image is correct.
"""

import os
import re
import unittest
import logging
import time
import json
import datetime
from parameterized import parameterized
import work_anniversary as work_anniv
from get_anniv_year_using_vqa import (
    get_anniversary_year,
    instantiate_model_and_processor,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    filename="anniversary_test.log",
    filemode="a",
)


def extract_anniv_year(image_path):
    "Extract the anniversary year from the image using VQA (OCR)"
    try:
        img_anniv_year = get_anniversary_year(
            image_path, TestWorkAnniversaryYear.model, TestWorkAnniversaryYear.processor
        )
        match = re.search(r"\d+", img_anniv_year)
        if match:
            return int(match.group())
        return None
    except Exception as error:
        logging.error("Failed to extract anniversary year: %s", error)
        return None


class TestWorkAnniversaryYear(unittest.TestCase):
    """
    Verify the anniversary year present in the work anniversary images
    """

    @classmethod
    def setUpClass(cls) -> None:
        cls.start_time = time.time()
        cls.model, cls.processor = instantiate_model_and_processor()
        cls.image_paths = cls.generate_test_images()

    @classmethod
    def generate_test_images(cls):
        "Prepares data and calls work anniversary lambda to generate a image"
        try:
            with open("employee_test_data.json") as json_file:
                data = json.load(json_file)
                employee_details = data["employees"]

            image_paths = []
            for employee in employee_details:
                emp_name = employee["emp_name"]
                emp_joined_date = datetime.datetime.strptime(
                    employee["joined_date"], "%Y-%m-%d"
                )
                current_date = datetime.datetime.strptime(
                    employee["current_date"], "%Y-%m-%d"
                )
                message, quote_string = work_anniv.calculate_work_anniversary(
                    emp_joined_date, current_date, emp_name
                )
                imagepath = work_anniv.add_text_to_image(
                    message, emp_name, quote_string
                )
                image_paths.append(imagepath)
            return image_paths
        except Exception as error:
            logging.info("Error generating test images: %s", error)

    @parameterized.expand(
        [
            (4, "Drishya TM_greeting_card.png"),
            (3, "Preedhi Vivek_greeting_card.png"),
        ]
    )
    def test_anniversary_year(self, anniversary_number, image_path):
        "Test the anniversary year extraction for different anniversary numbers"
        img_anniv_year = extract_anniv_year(image_path)
        self.assertEqual(
            img_anniv_year,
            anniversary_number,
            f"Expected {anniversary_number}, but got {img_anniv_year}",
        )

    @classmethod
    def tearDownClass(cls) -> None:
        "Remove existing images from test runs"
        for image_path in cls.image_paths:
            if os.path.exists(image_path):
                os.remove(image_path)
        elapsed_time = time.time() - cls.start_time
        logging.info("Test completed in %.2f seconds", elapsed_time)


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