Updated tests to match the new models.

tests/test_account.py

@@ -1,3 +1,4 @@
+import time
 import unittest
 from datetime import datetime
 from unittest.mock import patch
@@ -7,6 +8,7 @@ from pymongo.errors import ConnectionFailure
 
 from src.data.connection import ActionFailed, Collections, get_collection
 from src.data.repositories import account as account_repo
+from src.models.account import Account
 from src.utils.logger import logger
 from tests.base_test import BaseMongoTest
 
@@ -23,24 +25,20 @@ class TestAccountRepository(BaseMongoTest):
 	def test_init_functionality(self):
 		"""Test the init function's ability to create, drop, and preserve collections."""
 		self.assertIn(self.test_collection, self.db.list_collection_names())
-		# Test that data persists when drop=False
-		account_id = account_repo.create_account("Persist Test", "Doctor", collection_name=self.test_collection)
+		account_repo.create_account("Persist Test", "Doctor", collection_name=self.test_collection)
 		account_repo.init(collection_name=self.test_collection, drop=False)
 		self.assertEqual(get_collection(self.test_collection).count_documents({}), 1)
-		# Test that data is deleted when drop=True
 		account_repo.init(collection_name=self.test_collection, drop=True)
 		self.assertEqual(get_collection(self.test_collection).count_documents({}), 0)
 
 	def test_create_account(self):
 		"""Test successful account creation, including optional fields."""
-		# Test basic creation
 		name, role = "Test Doctor", "Doctor"
 		account_id = account_repo.create_account(name=name, role=role, collection_name=self.test_collection)
 		self.assertIsInstance(account_id, str)
 		doc = self.get_doc_by_id(Collections.ACCOUNT, account_id)
 		self.assertIsNotNone(doc)
 		self.assertEqual(doc["name"], name) # type: ignore
-		# Test creation with specialty
 		spec_id = account_repo.create_account("Spec", "Nurse", specialty="Cardiology", collection_name=self.test_collection)
 		spec_doc = self.get_doc_by_id(Collections.ACCOUNT, spec_id)
 		self.assertEqual(spec_doc["specialty"], "Cardiology") # type: ignore
@@ -50,8 +48,6 @@ class TestAccountRepository(BaseMongoTest):
 		account_id = account_repo.create_account("Update Logic", "Doctor", collection_name=self.test_collection)
 		original_doc = self.get_doc_by_id(Collections.ACCOUNT, account_id)
 		self.assertIsNotNone(original_doc)
-
-		# Test that 'created_at' is immutable
 		updates = {"name": "Updated Name", "created_at": datetime(2000, 1, 1)}
 		success = account_repo.update_account(account_id, updates, collection_name=self.test_collection)
 		self.assertTrue(success)
@@ -59,80 +55,56 @@ class TestAccountRepository(BaseMongoTest):
 		self.assertIsNotNone(updated_doc)
 		self.assertEqual(updated_doc["created_at"], original_doc["created_at"]) # type: ignore
 		self.assertLess(original_doc["updated_at"], updated_doc["updated_at"]) # type: ignore
-
-		# Test updating a non-existent account returns False
 		self.assertFalse(account_repo.update_account(str(ObjectId()), {"name": "No One"}, collection_name=self.test_collection))
 
 	def test_get_account_logic(self):
-		"""Test that get_account updates 'last_seen' but not 'updated_at'."""
+		"""Test that get_account updates 'last_seen' and returns a valid Account model."""
 		account_id = account_repo.create_account("GetMe", "Doctor", collection_name=self.test_collection)
 		original_doc = self.get_doc_by_id(Collections.ACCOUNT, account_id)
 		self.assertIsNotNone(original_doc)
-		self.assertNotIn("last_seen", original_doc) # type: ignore
+		time.sleep(0.01) # Ensure timestamp will be different
 
-		# Get the account, which should add 'last_seen'
 		account = account_repo.get_account(account_id, collection_name=self.test_collection)
 		self.assertIsNotNone(account)
-		self.assertIn("last_seen", account) # type: ignore
-
-		# Verify that 'updated_at' was NOT modified by the get operation
-		final_doc = self.get_doc_by_id(Collections.ACCOUNT, account_id)
-		self.assertIsNotNone(final_doc)
-		self.assertEqual(final_doc["updated_at"], original_doc["updated_at"]) # type: ignore
+		self.assertIsInstance(account, Account)
+		self.assertLess(original_doc["last_seen"], account.last_seen) # type: ignore
+		self.assertEqual(original_doc["updated_at"], account.updated_at) # type: ignore
+		self.assertIsNone(account_repo.get_account(str(ObjectId()), collection_name=self.test_collection))
 
 	def test_get_account_by_name(self):
 		"""Test retrieving an account by name and check for deprecation warning."""
 		name = "FindByName"
 		account_repo.create_account(name, "Nurse", collection_name=self.test_collection)
-
-		# Check that the function works and raises the expected warning
 		account = account_repo.get_account_by_name(name, collection_name=self.test_collection)
 		self.assertIsNotNone(account)
-		self.assertEqual(account["name"], name) # type: ignore
-
-		# Test retrieval of a non-existent name returns None
+		self.assertIsInstance(account, Account)
+		self.assertEqual(account.name, name) # type: ignore
 		self.assertIsNone(account_repo.get_account_by_name("NonExistent", collection_name=self.test_collection))
 
 	def test_search_accounts(self):
-		"""Test search functionality with various edge cases."""
+		"""Test search functionality returns a list of Account models."""
 		account_repo.create_account("Alpha Doctor", "Doctor", collection_name=self.test_collection)
 		account_repo.create_account("Beta Nurse", "Nurse", collection_name=self.test_collection)
-		account_repo.create_account("Charlie Doctor", "Doctor", collection_name=self.test_collection)
-
-		# Test case-insensitive partial match
-		self.assertEqual(len(account_repo.search_accounts("alpha", collection_name=self.test_collection)), 1)
-		# Test query matching multiple documents
-		self.assertEqual(len(account_repo.search_accounts("Doctor", collection_name=self.test_collection)), 2)
-		# Test limit parameter
-		self.assertEqual(len(account_repo.search_accounts("Doctor", limit=1, collection_name=self.test_collection)), 1)
-		# Test query with no matches
+		results = account_repo.search_accounts("alpha", collection_name=self.test_collection)
+		self.assertEqual(len(results), 1)
+		self.assertIsInstance(results[0], Account)
+		self.assertEqual(results[0].name, "Alpha Doctor")
 		self.assertEqual(len(account_repo.search_accounts("NonExistent", collection_name=self.test_collection)), 0)
-		# Test empty query string returns an empty list
-		self.assertEqual(len(account_repo.search_accounts("", collection_name=self.test_collection)), 0)
 
 	def test_get_all_accounts(self):
-		"""Test retrieving all accounts, verifying sorting and limit."""
+		"""Test retrieving all accounts, verifying sorting and model type."""
 		account_repo.create_account("Charlie", "Doctor", collection_name=self.test_collection)
 		account_repo.create_account("Alpha", "Nurse", collection_name=self.test_collection)
-		account_repo.create_account("Beta", "Caregiver", collection_name=self.test_collection)
-
 		all_accounts = account_repo.get_all_accounts(collection_name=self.test_collection)
-		self.assertEqual(len(all_accounts), 3)
-		# Verify results are sorted by name
-		self.assertEqual(all_accounts[0]["name"], "Alpha")
-		self.assertEqual(all_accounts[2]["name"], "Charlie")
-
-		# Test with a limit
-		limited_accounts = account_repo.get_all_accounts(limit=2, collection_name=self.test_collection)
-		self.assertEqual(len(limited_accounts), 2)
-		self.assertEqual(limited_accounts[1]["name"], "Beta")
+		self.assertEqual(len(all_accounts), 2)
+		self.assertIsInstance(all_accounts[0], Account)
+		self.assertEqual(all_accounts[0].name, "Alpha")
+		self.assertEqual(all_accounts[1].name, "Charlie")
 
 	def test_get_account_frame(self):
-		"""Test retrieving accounts as a pandas DataFrame for empty and populated collections."""
-		# Test with an empty collection
+		"""Test retrieving accounts as a pandas DataFrame."""
 		df_empty = account_repo.get_account_frame(collection_name=self.test_collection)
 		self.assertTrue(df_empty.empty)
-		# Test with data
 		account_repo.create_account("Frame Alpha", "Doctor", collection_name=self.test_collection)
 		df_full = account_repo.get_account_frame(collection_name=self.test_collection)
 		self.assertEqual(len(df_full), 1)
@@ -187,4 +159,4 @@ class TestAccountRepositoryExceptions(BaseMongoTest):
 if __name__ == "__main__":
 	logger().info("Starting MongoDB repository integration tests...")
 	unittest.main(verbosity=2)
-	logger().info("Tests completed and database connection closed.")
+	logger().info("Tests completed.")

tests/test_patient.py

@@ -6,6 +6,7 @@ from pymongo.errors import ConnectionFailure
 
 from src.data.connection import ActionFailed, Collections, get_collection
 from src.data.repositories import patient as patient_repo
+from src.models.patient import Patient
 from src.utils.logger import logger
 from tests.base_test import BaseMongoTest
 
@@ -27,7 +28,6 @@ class TestPatientRepository(BaseMongoTest):
 
 	def test_create_patient(self):
 		"""Test patient creation with minimal and full data."""
-		# Test minimal creation
 		patient_id = patient_repo.create_patient(
 			"John Doe", 45, "Male", "Caucasian", collection_name=self.test_collection
 		)
@@ -36,41 +36,26 @@ class TestPatientRepository(BaseMongoTest):
 		self.assertIsNotNone(doc)
 		self.assertEqual(doc["name"], "John Doe") # type: ignore
 
-		# Test full creation with all optional parameters
 		doctor_id = str(ObjectId())
 		full_id = patient_repo.create_patient(
-			name="Jane Doe",
-			age=30,
-			sex="Female",
-			ethnicity="Asian",
-			address="123 Wellness Way",
-			phone="555-123-4567",
-			email="jane.doe@example.com",
-			medications=["Lisinopril", "Metformin"],
-			past_assessment_summary="Routine check-up, generally healthy.",
-			assigned_doctor_id=doctor_id,
-			collection_name=self.test_collection
+			name="Jane Doe", age=30, sex="Female", ethnicity="Asian",
+			address="123 Wellness Way", phone="555-123-4567", email="jane.doe@example.com",
+			medications=["Lisinopril"], past_assessment_summary="Routine check-up.",
+			assigned_doctor_id=doctor_id, collection_name=self.test_collection
 		)
-		self.assertIsInstance(full_id, str)
 		full_doc = self.get_doc_by_id(Collections.PATIENT, full_id)
-
-		# Verify all fields were saved correctly
 		self.assertIsNotNone(full_doc)
-		self.assertEqual(full_doc["address"], "123 Wellness Way") # type: ignore
-		self.assertEqual(full_doc["phone"], "555-123-4567") # type: ignore
 		self.assertEqual(full_doc["email"], "jane.doe@example.com") # type: ignore
-		self.assertEqual(len(full_doc["medications"]), 2) # type: ignore
-		self.assertIn("Lisinopril", full_doc["medications"]) # type: ignore
-		self.assertEqual(full_doc["past_assessment_summary"], "Routine check-up, generally healthy.") # type: ignore
 		self.assertEqual(str(full_doc["assigned_doctor_id"]), doctor_id) # type: ignore
 
 	def test_get_patient_by_id(self):
-		"""Test retrieving an existing patient by ID."""
+		"""Test retrieving an existing patient by ID returns a Patient model."""
 		patient_id = patient_repo.create_patient("GetMe", 33, "Female", "Other", collection_name=self.test_collection)
 		patient = patient_repo.get_patient_by_id(patient_id, collection_name=self.test_collection)
 		self.assertIsNotNone(patient)
-		self.assertEqual(patient["_id"], patient_id) # type: ignore
-		# Test retrieval of a non-existent patient returns None
+		self.assertIsInstance(patient, Patient)
+		self.assertEqual(patient.id, patient_id) # type: ignore
+		self.assertEqual(patient.name, "GetMe") # type: ignore
 		self.assertIsNone(patient_repo.get_patient_by_id(str(ObjectId()), collection_name=self.test_collection))
 
 	def test_update_patient_profile(self):
@@ -80,18 +65,18 @@ class TestPatientRepository(BaseMongoTest):
 		modified_count = patient_repo.update_patient_profile(patient_id, updates, collection_name=self.test_collection)
 		self.assertEqual(modified_count, 1)
 		doc = self.get_doc_by_id(Collections.PATIENT, patient_id)
+		self.assertIsNotNone(doc)
 		self.assertEqual(doc["age"], 26) # type: ignore
-		# Test updating a non-existent patient returns a modified count of 0
 		self.assertEqual(patient_repo.update_patient_profile(str(ObjectId()), {"name": "Ghost"}, collection_name=self.test_collection), 0)
 
 	def test_search_patients(self):
-		"""Test patient search functionality with various queries and limits."""
+		"""Test patient search functionality returns a list of Patient models."""
 		patient_repo.create_patient("Alice Smith", 30, "Female", "Asian", collection_name=self.test_collection)
 		patient_repo.create_patient("Bob Smith", 45, "Male", "Caucasian", collection_name=self.test_collection)
 		results = patient_repo.search_patients("smith", collection_name=self.test_collection)
 		self.assertEqual(len(results), 2)
-		self.assertEqual(results[0]["name"], "Alice Smith")
-		# Test limit parameter
+		self.assertIsInstance(results[0], Patient)
+		self.assertEqual(results[0].name, "Alice Smith")
 		self.assertEqual(len(patient_repo.search_patients("s", limit=1, collection_name=self.test_collection)), 1)
 
 
@@ -106,11 +91,8 @@ class TestPatientRepositoryExceptions(BaseMongoTest):
 
 	def test_write_error_raises_action_failed(self):
 		"""Test that creating or updating with data violating schema raises ActionFailed."""
-		# Test creating a patient with a 'sex' value not in the schema's enum
 		with self.assertRaises(ActionFailed):
 			patient_repo.create_patient("Schema Test", 25, "InvalidValue", "Other", collection_name=self.test_collection)
-
-		# Test updating a patient with an invalid value
 		patient_id = patient_repo.create_patient("UpdateSchema", 30, "Male", "Other", collection_name=self.test_collection)
 		with self.assertRaises(ActionFailed):
 			patient_repo.update_patient_profile(patient_id, {"ethnicity": 123}, collection_name=self.test_collection)
@@ -128,7 +110,6 @@ class TestPatientRepositoryExceptions(BaseMongoTest):
 	def test_all_functions_raise_on_connection_error(self, mock_get_collection):
 		"""Test that all repo functions catch generic PyMongoErrors and raise ActionFailed."""
 		mock_get_collection.side_effect = ConnectionFailure("Simulated connection error")
-
 		with self.assertRaises(ActionFailed):
 			patient_repo.init(collection_name=self.test_collection, drop=True)
 		with self.assertRaises(ActionFailed):
@@ -143,4 +124,4 @@ class TestPatientRepositoryExceptions(BaseMongoTest):
 if __name__ == "__main__":
 	logger().info("Starting MongoDB repository integration tests...")
 	unittest.main(verbosity=2)
-	logger().info("Tests completed and database connection closed.")
+	logger().info("Tests completed.")

tests/test_session.py

@@ -8,6 +8,7 @@ from pymongo.errors import ConnectionFailure, WriteError
 
 from src.data.connection import ActionFailed, Collections, get_collection
 from src.data.repositories import session as session_repo
+from src.models.session import Message, Session
 from src.utils.logger import logger
 from tests.base_test import BaseMongoTest
 
@@ -34,46 +35,58 @@ class TestSessionRepository(BaseMongoTest):
 		session = session_repo.create_session(
 			self.account_id, self.patient_id, "Test Chat", collection_name=self.test_collection
 		)
-		self.assertIsInstance(session["_id"], str)
-		retrieved = session_repo.get_session(session["_id"], collection_name=self.test_collection)
+		self.assertIsInstance(session, Session)
+		self.assertEqual(session.title, "Test Chat")
+
+		retrieved = session_repo.get_session(session.id, collection_name=self.test_collection)
 		self.assertIsNotNone(retrieved)
-		self.assertEqual(retrieved["_id"], session["_id"]) # type: ignore
-		# Test getting a non-existent session
+		self.assertIsInstance(retrieved, Session)
+		self.assertEqual(retrieved.id, session.id) # type: ignore
+		self.assertEqual(retrieved.account_id, self.account_id) # type: ignore
 		self.assertIsNone(session_repo.get_session(str(ObjectId()), collection_name=self.test_collection))
 
 	def test_add_and_get_messages(self):
-		"""Test adding messages and retrieving them in the correct order."""
+		"""Test adding messages and retrieving them as Message models."""
 		session = session_repo.create_session(self.account_id, self.patient_id, "Msg Test", collection_name=self.test_collection)
-		session_repo.add_message(session["_id"], "User message 1", True, collection_name=self.test_collection)
-		session_repo.add_message(session["_id"], "AI response 1", False, collection_name=self.test_collection)
-		messages = session_repo.get_session_messages(session["_id"], collection_name=self.test_collection)
+
+		session_repo.add_message(session.id, "User message 1", True, collection_name=self.test_collection)
+		# Add a small delay to ensure a distinct timestamp for the next message
+		time.sleep(0.01)
+		session_repo.add_message(session.id, "AI response 1", False, collection_name=self.test_collection)
+
+		messages = session_repo.get_session_messages(session.id, collection_name=self.test_collection)
 		self.assertEqual(len(messages), 2)
-		self.assertEqual(messages[0]["content"], "AI response 1") # Descending order
-		self.assertEqual(messages[1]["_id"], 0)
-		# Test limit
-		self.assertEqual(len(session_repo.get_session_messages(session["_id"], limit=1, collection_name=self.test_collection)), 1)
+		self.assertIsInstance(messages[0], Message)
+		self.assertEqual(messages[0].content, "AI response 1") # Descending order is now guaranteed
+		self.assertEqual(messages[1].id, 0)
+		self.assertEqual(len(session_repo.get_session_messages(session.id, limit=1, collection_name=self.test_collection)), 1)
 
 	def test_list_sessions(self):
-		"""Test listing sessions for a specific patient and user, sorted by update time."""
+		"""Test listing sessions for a patient and user returns lists of Session models."""
 		session_repo.create_session(self.account_id, self.patient_id, "First", collection_name=self.test_collection)
-		time.sleep(0.01) # Ensure distinct timestamps
+		time.sleep(0.01)
 		s2 = session_repo.create_session(self.account_id, self.patient_id, "Second", collection_name=self.test_collection)
-		# Test listing for patient
+
 		patient_sessions = session_repo.list_patient_sessions(self.patient_id, collection_name=self.test_collection)
 		self.assertEqual(len(patient_sessions), 2)
-		self.assertEqual(patient_sessions[0]["_id"], s2["_id"]) # Most recent first
-		# Test listing for user
+		self.assertIsInstance(patient_sessions[0], Session)
+		self.assertEqual(patient_sessions[0].id, s2.id)
+
 		user_sessions = session_repo.get_user_sessions(self.account_id, collection_name=self.test_collection)
 		self.assertEqual(len(user_sessions), 2)
-		self.assertEqual(user_sessions[0]["_id"], s2["_id"])
+		self.assertEqual(user_sessions[0].id, s2.id)
 
 	def test_update_session_title(self):
 		"""Test updating a session's title and its timestamp."""
 		session = session_repo.create_session(self.account_id, self.patient_id, "Old", collection_name=self.test_collection)
-		original_doc = self.get_doc_by_id(Collections.SESSION, session["_id"])
-		success = session_repo.update_session_title(session["_id"], "New", collection_name=self.test_collection)
+		original_doc = self.get_doc_by_id(Collections.SESSION, session.id)
+		self.assertIsNotNone(original_doc)
+
+		success = session_repo.update_session_title(session.id, "New", collection_name=self.test_collection)
 		self.assertTrue(success)
-		updated_doc = self.get_doc_by_id(Collections.SESSION, session["_id"])
+
+		updated_doc = self.get_doc_by_id(Collections.SESSION, session.id)
+		self.assertIsNotNone(updated_doc)
 		self.assertEqual(updated_doc["title"], "New") # type: ignore
 		self.assertLess(original_doc["updated_at"], updated_doc["updated_at"]) # type: ignore
 		self.assertFalse(session_repo.update_session_title(str(ObjectId()), "Ghost", collection_name=self.test_collection))
@@ -81,27 +94,21 @@ class TestSessionRepository(BaseMongoTest):
 	def test_delete_session(self):
 		"""Test deleting a session."""
 		session = session_repo.create_session(self.account_id, self.patient_id, "To Delete", collection_name=self.test_collection)
-		self.assertTrue(session_repo.delete_session(session["_id"], collection_name=self.test_collection))
-		self.assertIsNone(session_repo.get_session(session["_id"], collection_name=self.test_collection))
+		self.assertTrue(session_repo.delete_session(session.id, collection_name=self.test_collection))
+		self.assertIsNone(session_repo.get_session(session.id, collection_name=self.test_collection))
 		self.assertFalse(session_repo.delete_session(str(ObjectId()), collection_name=self.test_collection))
 
 	def test_prune_old_sessions(self):
 		"""Test deleting sessions older than a specified number of days."""
-		# Create two valid sessions
-		old_session = session_repo.create_session(self.account_id, self.patient_id, "Old Session", collection_name=self.test_collection)
-		session_repo.create_session(self.account_id, self.patient_id, "New Session", collection_name=self.test_collection)
-
-		# Manually update one session to be old
+		old_session = session_repo.create_session(self.account_id, self.patient_id, "Old", collection_name=self.test_collection)
+		session_repo.create_session(self.account_id, self.patient_id, "New", collection_name=self.test_collection)
 		old_date = datetime.now(timezone.utc) - timedelta(days=31)
 		get_collection(self.test_collection).update_one(
-			{"_id": ObjectId(old_session["_id"])},
-			{"$set": {"updated_at": old_date}}
+			{"_id": ObjectId(old_session.id)}, {"$set": {"updated_at": old_date}}
 		)
-
 		self.assertEqual(get_collection(self.test_collection).count_documents({}), 2)
 		deleted_count = session_repo.prune_old_sessions(days=30, collection_name=self.test_collection)
 		self.assertEqual(deleted_count, 1)
-		self.assertEqual(get_collection(self.test_collection).count_documents({}), 1)
 
 
 class TestSessionRepositoryExceptions(BaseMongoTest):
@@ -118,14 +125,11 @@ class TestSessionRepositoryExceptions(BaseMongoTest):
 	@patch('src.data.repositories.session.get_collection')
 	def test_write_error_raises_action_failed(self, mock_get_collection):
 		"""Test that a WriteError during an operation is raised as ActionFailed."""
-		# Configure the mock to return a collection object whose methods raise errors
 		mock_collection = mock_get_collection.return_value
 		mock_collection.update_one.side_effect = WriteError("Simulated schema validation error")
-		mock_collection.find_one.return_value = {"messages": []} # Needed for add_message to proceed
-
+		mock_collection.find_one.return_value = {"messages": []}
 		with self.assertRaises(ActionFailed):
-			# This will fail inside at the update_one call, which we've mocked
-			session_repo.add_message("68e212480769b3f99015f43c", "content", True, collection_name=self.test_collection)
+			session_repo.add_message(str(ObjectId()), "content", True, collection_name=self.test_collection)
 
 	def test_invalid_id_raises_action_failed(self):
 		"""Test that functions raise ActionFailed when given a malformed ObjectId string."""
@@ -150,7 +154,6 @@ class TestSessionRepositoryExceptions(BaseMongoTest):
 	def test_all_functions_raise_on_connection_error(self, mock_get_collection):
 		"""Test that all repo functions catch generic PyMongoErrors and raise ActionFailed."""
 		mock_get_collection.side_effect = ConnectionFailure("Simulated connection error")
-
 		with self.assertRaises(ActionFailed):
 			session_repo.init(collection_name=self.test_collection, drop=True)
 		with self.assertRaises(ActionFailed):
@@ -175,4 +178,4 @@ class TestSessionRepositoryExceptions(BaseMongoTest):
 if __name__ == "__main__":
 	logger().info("Starting MongoDB repository integration tests...")
 	unittest.main(verbosity=2)
-	logger().info("Tests completed and database connection closed.")
+	logger().info("Tests completed.")