import unittest from tempfile import NamedTemporaryFile from shop_ledger.processor import extract_document_text, prepare_document_input from shop_ledger.ui import ( add_to_ledger, apply_row_correction, ask_ledger, ask_ledger_chat, ask_ledger_voice_chat, choose_input, compose_chart, generate_daily_brief, initial_ask_chat, run_command_palette, ) class InputChoiceTests(unittest.TestCase): def test_auto_asks_when_text_and_audio_exist(self): choice = choose_input("paid Ravi 1200", "/tmp/audio.wav", None, "Auto") self.assertEqual(choice["status"], "conflict") self.assertIn("Multiple inputs", choice["notice"]) def test_text_choice_uses_text_when_audio_exists(self): choice = choose_input("paid Ravi 1200", "/tmp/audio.wav", None, "Text note") self.assertEqual(choice["status"], "ready") self.assertEqual(choice["source"], "text") def test_auto_uses_audio_when_audio_is_only_input(self): choice = choose_input("", "/tmp/audio.wav", None, "Auto") self.assertEqual(choice["status"], "ready") self.assertEqual(choice["source"], "audio") def test_auto_uses_document_when_document_is_only_input(self): choice = choose_input("", None, "/tmp/receipt.pdf", "Auto") self.assertEqual(choice["status"], "ready") self.assertEqual(choice["source"], "document") def test_document_text_extraction_reads_plain_text_files(self): with NamedTemporaryFile("w", suffix=".txt") as handle: handle.write("paid Ravi 1200 for rice bags") handle.flush() text = extract_document_text(handle.name) self.assertIn("Ravi", text) def test_document_image_preparation_creates_data_url(self): from PIL import Image with NamedTemporaryFile(suffix=".png") as handle: Image.new("RGB", (8, 8), color="white").save(handle.name) document = prepare_document_input(handle.name) self.assertEqual(document["kind"], "image") self.assertTrue(document["image_urls"][0].startswith("data:image/jpeg;base64,")) def test_successful_text_add_clears_written_note(self): def fake_process(note, currency, image_urls=None): return { "entries": [ { "date": "2026-06-11", "direction": "expense", "counterparty": "Ravi", "item": "rice bags", "quantity": "", "amount": 1200, "currency": currency, "category": "inventory", "payment_status": "paid", "due_date": "", "confidence": 0.9, "reminder": "", } ], "reminders": [], "questions": [], "model_used": "fake", } output = add_to_ledger("paid Ravi 1200", None, None, "Auto", "LKR", [], fake_process) self.assertEqual(len(output[6]), 1) self.assertEqual(output[7]["value"], "") self.assertEqual(output[10]["value"], "Auto") self.assertIn("Added 1 row", output[11]) def test_successful_document_add_sends_image_urls_and_clears_file(self): captured = {} def fake_process(note, currency, image_urls=None): captured["note"] = note captured["image_urls"] = image_urls return { "entries": [ { "date": "2026-06-11", "direction": "expense", "counterparty": "Ravi", "item": "rice bags", "quantity": "", "amount": 1200, "currency": currency, "category": "inventory", "payment_status": "paid", "due_date": "", "confidence": 0.9, "reminder": "", } ], "reminders": [], "questions": [], "model_used": "fake", } with NamedTemporaryFile("w", suffix=".txt") as handle: handle.write("paid Ravi 1200 for rice bags") handle.flush() output = add_to_ledger("", None, handle.name, "Document", "LKR", [], fake_process) self.assertIn("paid Ravi", captured["note"]) self.assertIsNone(captured["image_urls"]) self.assertEqual(output[9]["value"], None) self.assertIn("Added 1 row", output[11]) def test_generate_daily_brief_uses_supplied_function(self): rows = [{"amount": 1200, "currency": "LKR", "direction": "expense", "payment_status": "paid"}] markdown = generate_daily_brief( rows, "LKR", lambda supplied_rows, currency: {"brief": f"{len(supplied_rows)} rows in {currency}", "model_used": "fake"}, ) self.assertIn("1 rows in LKR", markdown) self.assertIn("fake", markdown) def test_ask_ledger_uses_supplied_function(self): rows = [{"amount": 7500, "currency": "LKR", "payment_status": "due"}] markdown = ask_ledger( rows, "Who owes me most?", "LKR", lambda supplied_rows, question, currency: {"answer": f"{question} / {len(supplied_rows)}", "model_used": "fake"}, ) self.assertIn("Who owes me most?", markdown) self.assertIn("fake", markdown) def test_ask_ledger_chat_appends_messages_and_clears_input(self): rows = [{"amount": 7500, "currency": "LKR", "payment_status": "due"}] history, next_question = ask_ledger_chat( rows, "Who owes me most?", initial_ask_chat(), "LKR", lambda supplied_rows, question, currency: {"answer": "Nimal owes LKR 7,500.", "model_used": "fake"}, ) self.assertEqual(next_question, "") self.assertEqual(history[-2]["role"], "user") self.assertEqual(history[-1]["role"], "assistant") self.assertIn("Nimal", history[-1]["content"]) def test_ask_ledger_voice_chat_transcribes_and_answers(self): history, next_question, next_audio = ask_ledger_voice_chat( [{"counterparty": "Nimal", "amount": 7500, "payment_status": "due", "currency": "LKR"}], "/tmp/question.wav", initial_ask_chat(), "LKR", lambda rows, question, currency: {"answer": f"Answered: {question}", "model_used": "fake"}, transcribe_fn=lambda path: "Who owes me most?", ) self.assertEqual(next_question, "") self.assertIsNone(next_audio) self.assertIn("Who owes me most?", history[-2]["content"]) self.assertIn("Answered", history[-1]["content"]) def test_ask_ledger_voice_chat_handles_empty_transcript(self): history, _, next_audio = ask_ledger_voice_chat( [], "/tmp/question.wav", initial_ask_chat(), "LKR", lambda rows, question, currency: {"answer": "unused", "model_used": "fake"}, transcribe_fn=lambda path: "", ) self.assertIsNone(next_audio) self.assertIn("could not hear", history[-1]["content"]) def test_run_command_palette_uses_current_rows(self): rows = [{"payment_status": "due", "counterparty": "Nimal", "amount": 7500, "currency": "LKR", "item": "tea"}] output = run_command_palette(rows, "Show unpaid") self.assertIn("Nimal", output) def test_compose_chart_returns_markdown_figure_and_clears_input(self): rows = [{"payment_status": "due", "counterparty": "Nimal", "amount": 7500, "currency": "LKR"}] markdown, figure, next_question = compose_chart( rows, "Who owes me?", lambda supplied_rows, question: {"chart": "due_by_party", "reason": "Dues", "model_used": "fake"}, ) self.assertIn("Due radar", markdown) self.assertTrue(hasattr(figure, "to_plotly_json")) self.assertEqual(next_question, "") def test_apply_row_correction_updates_state_and_confidence(self): rows = [ { "date": "2026-06-11", "direction": "income", "counterparty": "", "item": "tea packets", "quantity": "", "amount": 750, "currency": "LKR", "category": "sales", "payment_status": "due", "due_date": "", "confidence": 0.42, "reminder": "", } ] output = apply_row_correction(rows, 1, "counterparty", "Nimal", "LKR") updated_rows = output[6] self.assertEqual(updated_rows[0]["counterparty"], "Nimal") self.assertEqual(updated_rows[0]["confidence"], 0.9) self.assertIn("Updated row 1", output[-1]) def test_apply_row_correction_rejects_bad_amount(self): rows = [ { "date": "2026-06-11", "direction": "income", "counterparty": "Nimal", "item": "tea packets", "quantity": "", "amount": 750, "currency": "LKR", "category": "sales", "payment_status": "due", "due_date": "", "confidence": 0.42, "reminder": "", } ] output = apply_row_correction(rows, 1, "amount", "many rupees", "LKR") self.assertEqual(output[6][0]["amount"], 750) self.assertIn("need a number", output[-1]) if __name__ == "__main__": unittest.main()