gauravmeena0708/epfo-circulars

1import unittest2import pandas as pd3import numpy as np4 5from data_assistant import (6    DataExtractionError,7    load_csv_dataframe,8    generate_dataset_profile,9    search_dataframe,10    prepare_dataframe_llm_context,11)12 13 14class DataAssistantTests(unittest.TestCase):15    def setUp(self):16        self.sample_csv_text = """Issue_ID,Issue_Topic,Category,Functional_Owner,Status,Count_Affected17101,Form 13 Transfer Request stuck at employer,Transfers,Field Office Delhi,Open,1518102,DSC registration error on Employer Portal,DSC,Tech Backend,In Progress,819103,Form 13 Member Passbook not updating after transfer,Transfers,Tech Backend,Open,2420104,Joint Declaration name mismatch,Member Profile,Field Office Mumbai,Resolved,521105,Form 19 Claim rejection due to KYC,Claims,Settlement Team,Open,1222106,Form 13 Annexure K generation failed,Transfers,Tech Backend,Resolved,323"""24        self.sample_csv_bytes = self.sample_csv_text.encode("utf-8")25 26    def test_load_valid_csv(self):27        df = load_csv_dataframe(self.sample_csv_bytes, "cites_sample.csv")28        self.assertEqual(len(df), 6)29        self.assertIn("Issue_Topic", df.columns)30        self.assertIn("Functional_Owner", df.columns)31 32    def test_arbitrary_schema_and_double_csv_suffix_are_supported(self):33        location_csv = (34            b"place,latitude,longitude,population\n"35            b"North Point,28.61,77.21,1200\n"36            b"South Point,19.08,72.88,900\n"37        )38        df = load_csv_dataframe(39            location_csv,40            "~/Downloads/location.csv - location.csv.csv",41        )42 43        self.assertEqual(list(df.columns), ["place", "latitude", "longitude", "population"])44        self.assertEqual(len(df), 2)45 46    def test_domain_words_are_not_discarded_from_deterministic_search(self):47        df = load_csv_dataframe(b"category\nissues\nother\n", "generic.csv")48        context = prepare_dataframe_llm_context(df, "How many issues are there?")49 50        self.assertIn("Exactly **1** matching row(s) found", context)51 52    def test_load_empty_csv_raises_error(self):53        with self.assertRaises(DataExtractionError):54            load_csv_dataframe(b"", "empty.csv")55        with self.assertRaises(DataExtractionError):56            load_csv_dataframe(b"   \n  ", "whitespace.csv")57 58    def test_load_non_csv_and_malformed_rows_raise_error(self):59        with self.assertRaises(DataExtractionError):60            load_csv_dataframe(b"not actually csv", "plain.txt")61        with self.assertRaises(DataExtractionError):62            load_csv_dataframe(b"a,b\n1,2\n3,4,5\n6,7\n", "broken.csv")63 64    def test_duplicate_trimmed_columns_are_made_unique(self):65        df = load_csv_dataframe(b"Name, Name\nAlice,Bob\n", "duplicates.csv")66        self.assertEqual(list(df.columns), ["Name", "Name_2"])67        profile = generate_dataset_profile(df)68        self.assertEqual(profile.column_count, 2)69 70    def test_load_semicolon_delimited_csv(self):71        semicolon_csv = b"ColA;ColB;ColC\n1;2;3\n4;5;6\n"72        df = load_csv_dataframe(semicolon_csv, "semi.csv")73        self.assertEqual(len(df), 2)74        self.assertEqual(len(df.columns), 3)75        self.assertIn("ColA", df.columns)76 77    def test_dataset_profile_generation(self):78        df = load_csv_dataframe(self.sample_csv_bytes)79        profile = generate_dataset_profile(df)80        self.assertEqual(profile.row_count, 6)81        self.assertEqual(profile.column_count, 6)82        self.assertIn("Count_Affected", profile.numeric_columns)83        self.assertIn("Issue_Topic", profile.categorical_columns)84        self.assertIn("Count_Affected", profile.summary_stats)85        self.assertEqual(profile.summary_stats["Count_Affected"]["min"], 3.0)86        self.assertEqual(profile.summary_stats["Count_Affected"]["max"], 24.0)87 88    def test_search_dataframe_form_13_exact_count(self):89        df = load_csv_dataframe(self.sample_csv_bytes)90        matched_df, total_matches, breakdown = search_dataframe(df, "Form 13")91        92        # In our sample CSV: rows 101, 103, 106 have "Form 13" in Issue_Topic93        self.assertEqual(total_matches, 3)94        self.assertEqual(len(matched_df), 3)95        self.assertEqual(breakdown.get("Issue_Topic"), 3)96 97    def test_search_dataframe_case_insensitive(self):98        df = load_csv_dataframe(self.sample_csv_bytes)99        matched_df, total_matches, _ = search_dataframe(df, "form 13", case_sensitive=False)100        self.assertEqual(total_matches, 3)101 102    def test_search_dataframe_owner_filter(self):103        df = load_csv_dataframe(self.sample_csv_bytes)104        matched_df, total_matches, _ = search_dataframe(df, "Tech Backend")105        self.assertEqual(total_matches, 3)106 107    def test_prepare_dataframe_llm_context_includes_deterministic_count(self):108        df = load_csv_dataframe(self.sample_csv_bytes)109        user_query = "How many Form 13 issues are there in the dataset?"110        context = prepare_dataframe_llm_context(df, user_query)111        112        self.assertIn("Dataset Structure", context)113        self.assertIn("Form 13", context)114        self.assertIn("Exactly **3** matching row(s) found", context)115 116    def test_prepare_context_includes_combined_filter_count(self):117        df = load_csv_dataframe(self.sample_csv_bytes)118        context = prepare_dataframe_llm_context(119            df,120            "How many Form 13 issues are Open?",121        )122 123        self.assertIn("Combined filter", context)124        self.assertIn("Exactly **2** row(s) match all identified terms", context)125 126    def test_prepare_context_respects_character_limit(self):127        df = pd.DataFrame(128            {f"Column_{index}": ["x" * 200 for _ in range(20)] for index in range(40)}129        )130        context = prepare_dataframe_llm_context(131            df,132            "summarize this dataset",133            max_context_chars=2_500,134        )135 136        self.assertLessEqual(len(context), 2_500)137        self.assertIn("dataset context truncated", context)138 139 140if __name__ == "__main__":141    unittest.main()142