doczyai-pipelines/src/tests/test_fieldset.py

import json
import os
import unittest
from unittest.mock import mock_open, patch

from src.constants.constants import Constants
from src.prompts.fieldset import Field, FieldSet, _cache_lock, _field_data_cache


class TestFieldWithRealConstants(unittest.TestCase):
    """Tests for the Field class using real Constants values."""

    def setUp(self):
        """Set up test fixtures before each test method."""
        # Create test field with a real VALID_BILL_TYPE placeholder
        self.test_field_dict = {
            "field_name": "TEST_FIELD",
            "relationship": "one-to-one",
            "field_type": "string",
            "prompt": "What is the test field? Valid values: {valid_values}",
            "valid_values": "VALID_BILL_TYPE",
            "base_field": "base_test_field",
            "crosswalk": "test_crosswalk",
            "keywords": ["test", "field"],
            "methodology": "LLM",
            "case_sensitive": True,
            "regex_pattern": r"\d+",
            "format": "text",
            "allow_na": False,
        }
        self.field = Field(self.test_field_dict)

        # Create a real Constants object - this will load actual mappings
        self.constants = Constants()

    def test_init(self):
        """Test Field initialization with all parameters."""
        # Test all fields are correctly set
        self.assertEqual(self.field.field_name, "TEST_FIELD")
        self.assertEqual(self.field.relationship, "one-to-one")
        self.assertEqual(self.field.field_type, "string")
        self.assertEqual(
            self.field.prompt, "What is the test field? Valid values: {valid_values}"
        )
        self.assertEqual(self.field.valid_values, "VALID_BILL_TYPE")
        self.assertEqual(self.field.base_field, "base_test_field")
        self.assertEqual(self.field.crosswalk, "test_crosswalk")
        self.assertEqual(self.field.keywords, ["test", "field"])
        self.assertEqual(self.field.methodology, "LLM")
        self.assertTrue(self.field.case_sensitive)
        self.assertEqual(self.field.regex_pattern, r"\d+")
        self.assertEqual(self.field.format, "text")
        self.assertFalse(self.field.allow_na)
        self.assertIsNone(self.field.retrieved_text)

    def test_init_minimal(self):
        """Test Field initialization with minimal parameters."""
        minimal_field_dict = {
            "field_name": "MINIMAL_FIELD",
            "relationship": "one-to-one",
            "field_type": "string",
        }
        field = Field(minimal_field_dict)

        # Check required fields
        self.assertEqual(field.field_name, "MINIMAL_FIELD")
        self.assertEqual(field.relationship, "one-to-one")
        self.assertEqual(field.field_type, "string")

        # Check defaults for optional fields
        self.assertIsNone(field.prompt)
        self.assertIsNone(field.valid_values)
        self.assertIsNone(field.base_field)
        self.assertIsNone(field.crosswalk)
        self.assertIsNone(field.keywords)
        self.assertIsNone(field.methodology)
        self.assertFalse(field.case_sensitive)
        self.assertIsNone(field.regex_pattern)
        self.assertIsNone(field.format)
        self.assertTrue(field.allow_na)

    def test_repr(self):
        """Test the string representation of a Field object."""
        self.assertEqual(repr(self.field), "TEST_FIELD")

    @patch(
        "builtins.open",
        new_callable=mock_open,
        read_data=json.dumps(
            [
                {
                    "field_name": "FIELD1",
                    "relationship": "one-to-one",
                    "field_type": "string",
                },
                {
                    "field_name": "FIELD2",
                    "relationship": "one-to-many",
                    "field_type": "number",
                },
            ]
        ),
    )
    def test_load_from_file(self, mock_file):
        """Test loading a Field object from a file."""
        field = Field.load_from_file("dummy_path.json", "FIELD2")
        self.assertEqual(field.field_name, "FIELD2")
        self.assertEqual(field.relationship, "one-to-many")
        self.assertEqual(field.field_type, "number")

    def test_from_values(self):
        """Test creating a Field using the from_values class method."""
        field = Field.from_values(
            field_name="TEST_FIELD2",
            relationship="one-to-many",
            field_type="number",
            prompt="Test prompt",
        )
        self.assertEqual(field.field_name, "TEST_FIELD2")
        self.assertEqual(field.relationship, "one-to-many")
        self.assertEqual(field.field_type, "number")
        self.assertEqual(field.prompt, "Test prompt")

    def test_to_dict(self):
        """Test converting a Field to a dictionary."""
        field_dict = self.field.to_dict()
        self.assertEqual(field_dict["field_name"], "TEST_FIELD")
        self.assertEqual(field_dict["relationship"], "one-to-one")
        self.assertEqual(field_dict["field_type"], "string")
        self.assertEqual(
            field_dict["prompt"], "What is the test field? Valid values: {valid_values}"
        )
        self.assertEqual(field_dict["valid_values"], "VALID_BILL_TYPE")
        self.assertEqual(field_dict["base_field"], "base_test_field")
        self.assertEqual(field_dict["keywords"], ["test", "field"])
        self.assertEqual(field_dict["methodology"], "LLM")
        self.assertTrue(field_dict["case_sensitive"])
        self.assertEqual(field_dict["regex_pattern"], r"\d+")
        self.assertEqual(field_dict["format"], "text")

    @patch("builtins.print")
    def test_print(self, mock_print):
        """Test printing a Field object."""
        self.field.print()
        # Verify print was called the expected number of times
        self.assertEqual(mock_print.call_count, 9)

    def test_get_prompt_dict_with_real_constants(self):
        """Test getting a prompt dictionary from a Field with real Constants."""
        prompt_dict = self.field.get_prompt_dict(self.constants)

        # Get the actual list of bill types from Constants
        bill_types = self.constants.VALID_BILL_TYPE

        # Verify the prompt contains the actual bill types from Constants
        prompt = prompt_dict["TEST_FIELD"]
        self.assertTrue(prompt.startswith("What is the test field? Valid values: "))

        # Check if the output contains the real values
        for bill_type in bill_types[:3]:  # Check first few to avoid overly long test
            self.assertIn(bill_type, prompt)

    def test_get_prompt_no_valid_values(self):
        """Test getting a prompt with no valid_values to resolve."""
        field = Field(
            {
                "field_name": "SIMPLE_FIELD",
                "relationship": "one-to-one",
                "field_type": "string",
                "prompt": "What is the simple field?",
            }
        )
        prompt = field.get_prompt()
        self.assertEqual(prompt, "What is the simple field?")

    def test_get_prompt_with_real_constants(self):
        """Test resolving valid values in a prompt using real Constants."""
        prompt = self.field.get_prompt(self.constants)

        # Check that the prompt includes the actual bill types
        self.assertTrue(prompt.startswith("What is the test field? Valid values: "))

        # Verify that the real constants values are used
        bill_types = self.constants.VALID_BILL_TYPE
        bill_types_str = str(bill_types)
        self.assertIn(bill_types_str, prompt)

    def test_update_valid_values(self):
        """Test updating valid values for a Field."""
        self.field.update_valid_values("NEW_VALID_VALUES")
        self.assertEqual(self.field.valid_values, "NEW_VALID_VALUES")


class TestFieldSetWithRealConstants(unittest.TestCase):
    """Tests for the FieldSet class using real Constants values."""

    def setUp(self):
        """Set up test fixtures before each test method."""
        # Create fields with real Constants placeholders
        self.test_fields = [
            Field(
                {
                    "field_name": "FIELD1",
                    "relationship": "one-to-one",
                    "field_type": "string",
                    "prompt": "What is field 1?",
                }
            ),
            Field(
                {
                    "field_name": "FIELD2",
                    "relationship": "one-to-many",
                    "field_type": "number",
                    "prompt": "What is field 2? Valid values: {valid_values}",
                    "valid_values": "VALID_BILL_TYPE",
                }
            ),
            Field(
                {
                    "field_name": "FIELD3",
                    "relationship": "one-to-one",
                    "field_type": "string",
                    "prompt": "What is field 3? Valid values: {valid_values}",
                    "valid_values": "VALID_LOBS",
                }
            ),
        ]
        self.fieldset = FieldSet(fields=self.test_fields)

        # Create a real Constants object
        self.constants = Constants()

    @patch(
        "builtins.open",
        new_callable=mock_open,
        read_data=json.dumps(
            [
                {
                    "field_name": "FIELD1",
                    "relationship": "one-to-one",
                    "field_type": "string",
                },
                {
                    "field_name": "FIELD2",
                    "relationship": "one-to-many",
                    "field_type": "number",
                },
                {
                    "field_name": "FIELD3",
                    "relationship": "one-to-one",
                    "field_type": "boolean",
                },
            ]
        ),
    )
    def test_init_with_file(self, mock_file):
        """Test FieldSet initialization with a file."""

        # Clear the cache so that our mock will be used
        with _cache_lock:
            _field_data_cache.clear()
        fieldset = FieldSet(file_path="dummy_path.json")
        self.assertEqual(len(fieldset.fields), 3)
        self.assertEqual(fieldset.fields[0].field_name, "FIELD1")
        self.assertEqual(fieldset.fields[1].field_name, "FIELD2")
        self.assertEqual(fieldset.fields[2].field_name, "FIELD3")

    def test_add_field(self):
        """Test adding a field to a FieldSet."""
        new_field = Field(
            {
                "field_name": "FIELD4",
                "relationship": "one-to-one",
                "field_type": "string",
            }
        )
        self.fieldset.add_field(new_field)
        self.assertEqual(len(self.fieldset.fields), 4)
        self.assertEqual(self.fieldset.fields[3].field_name, "FIELD4")

        # Test adding a duplicate field (should not add)
        duplicate_field = Field(
            {
                "field_name": "FIELD1",
                "relationship": "one-to-one",
                "field_type": "string",
            }
        )
        self.fieldset.add_field(duplicate_field)
        self.assertEqual(len(self.fieldset.fields), 4)  # No change in length

    def test_get_field(self):
        """Test getting a field by name."""
        field = self.fieldset.get_field("FIELD2")
        self.assertEqual(field.field_name, "FIELD2")
        self.assertEqual(field.relationship, "one-to-many")

        # Test getting a field that doesn't exist
        with self.assertRaises(ValueError):
            self.fieldset.get_field("NONEXISTENT_FIELD")

    def test_list_fields(self):
        """Test listing all field names."""
        field_names = self.fieldset.list_fields()
        self.assertEqual(field_names, ["FIELD1", "FIELD2", "FIELD3"])

    def test_remove_field(self):
        """Test removing a field."""
        self.fieldset.remove_field("FIELD2")
        self.assertEqual(len(self.fieldset.fields), 2)
        self.assertEqual(self.fieldset.fields[0].field_name, "FIELD1")
        self.assertEqual(self.fieldset.fields[1].field_name, "FIELD3")

    def test_update_field_attr(self):
        """Test updating a field attribute."""
        self.fieldset.update_field_attr("FIELD1", "prompt", "Updated prompt")
        self.assertEqual(self.fieldset.get_field("FIELD1").prompt, "Updated prompt")

    def test_get_prompt_dict_with_real_constants(self):
        """Test getting prompt dictionaries with real Constants values."""
        prompt_dict = self.fieldset.get_prompt_dict(self.constants)
        self.assertEqual(len(prompt_dict), 3)
        self.assertEqual(prompt_dict["FIELD1"], "What is field 1?")

        # Check that FIELD2 has the actual bill types
        self.assertTrue(
            prompt_dict["FIELD2"].startswith("What is field 2? Valid values: ")
        )
        for bill_type in self.constants.VALID_BILL_TYPE[:3]:  # Check first few
            self.assertIn(bill_type, str(prompt_dict["FIELD2"]))

        # Check that FIELD3 has the actual LOBs
        self.assertTrue(
            prompt_dict["FIELD3"].startswith("What is field 3? Valid values: ")
        )
        for lob in self.constants.VALID_LOBS[:3]:  # Check first few
            self.assertIn(lob, str(prompt_dict["FIELD3"]))

        # Test filtering by field_type
        prompt_dict = self.fieldset.get_prompt_dict(self.constants, field_type="string")
        self.assertEqual(len(prompt_dict), 2)
        self.assertEqual(prompt_dict["FIELD1"], "What is field 1?")
        self.assertIn("FIELD3", prompt_dict)

        # Test filtering by relationship
        prompt_dict = self.fieldset.get_prompt_dict(
            self.constants, relationship="one-to-many"
        )
        self.assertEqual(len(prompt_dict), 1)
        self.assertIn("FIELD2", prompt_dict)

    def test_print_prompt_dict_with_real_constants(self):
        """Test printing all prompts as a string with real Constants."""
        prompt_str = self.fieldset.print_prompt_dict(self.constants)

        # Verify each field is in the output
        self.assertIn("FIELD1 : What is field 1?", prompt_str)
        self.assertIn("FIELD2 : What is field 2? Valid values: ", prompt_str)
        self.assertIn("FIELD3 : What is field 3? Valid values: ", prompt_str)

        # Verify real constants values are used
        for bill_type in self.constants.VALID_BILL_TYPE[:2]:  # Check first couple
            self.assertIn(bill_type, prompt_str)

        for lob in self.constants.VALID_LOBS[:2]:  # Check first couple
            self.assertIn(lob, prompt_str)

    @patch("builtins.print")
    def test_print(self, mock_print):
        """Test printing all fields."""
        self.fieldset.print()
        # Each field calls print multiple times
        self.assertTrue(mock_print.call_count > 20)

    def test_repr(self):
        """Test the string representation of a FieldSet."""
        self.assertEqual(repr(self.fieldset), "['FIELD1', 'FIELD2', 'FIELD3']")

        # Test empty FieldSet
        empty_fieldset = FieldSet()
        self.assertEqual(repr(empty_fieldset), "FieldSet: No fields loaded")

    @patch("builtins.print")
    def test_print_fields(self, mock_print):
        """Test printing field names."""
        self.fieldset.print_fields()
        mock_print.assert_called_with(["FIELD1", "FIELD2", "FIELD3"])

    def test_filter(self):
        """Test filtering fields based on criteria."""
        # Filter by relationship
        filtered = self.fieldset.filter(relationship="one-to-one")
        self.assertEqual(len(filtered.fields), 2)
        self.assertEqual(filtered.fields[0].field_name, "FIELD1")
        self.assertEqual(filtered.fields[1].field_name, "FIELD3")

        # Filter by field_type
        filtered = self.fieldset.filter(field_type="number")
        self.assertEqual(len(filtered.fields), 1)
        self.assertEqual(filtered.fields[0].field_name, "FIELD2")

    def test_contains_fields(self):
        """Test checking if FieldSet contains fields."""
        self.assertTrue(self.fieldset.contains_fields())

        # Test empty FieldSet
        empty_fieldset = FieldSet()
        self.assertFalse(empty_fieldset.contains_fields())

    def test_combine(self):
        """Test combining FieldSets."""
        other_fields = [
            Field(
                {
                    "field_name": "FIELD4",
                    "relationship": "one-to-one",
                    "field_type": "string",
                }
            ),
            Field(
                {
                    "field_name": "FIELD1",  # Duplicate, should not be added
                    "relationship": "one-to-one",
                    "field_type": "string",
                }
            ),
        ]
        other_fieldset = FieldSet(fields=other_fields)

        # Test combine with new FieldSet (not in-place)
        combined = self.fieldset.combine(other_fieldset)
        self.assertEqual(len(combined.fields), 4)
        self.assertEqual(combined.fields[0].field_name, "FIELD1")
        self.assertEqual(combined.fields[1].field_name, "FIELD2")
        self.assertEqual(combined.fields[2].field_name, "FIELD3")
        self.assertEqual(combined.fields[3].field_name, "FIELD4")

        # Original should be unchanged
        self.assertEqual(len(self.fieldset.fields), 3)

        # Test combine in-place
        self.fieldset.combine(other_fieldset, inplace=True)
        self.assertEqual(len(self.fieldset.fields), 4)
        self.assertEqual(self.fieldset.fields[3].field_name, "FIELD4")


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