Updated session tests

tests/test_session.py

@@ -1,8 +1,10 @@
 import time
 import unittest
 from datetime import datetime, timedelta, timezone
+from unittest.mock import patch
 
 from bson import ObjectId
+from pymongo.errors import ConnectionFailure, WriteError
 
 from src.data.connection import (ActionFailed, Collections, close_connection,
                                  get_collection)
@@ -12,13 +14,13 @@ from tests.base_test import BaseMongoTest
 
 
 class TestSessionRepository(BaseMongoTest):
+	"""Test class for the 'happy path' and edge cases of session repository functions."""
 
 	def setUp(self):
 		"""Set up a clean test environment before each test."""
 		super().setUp()
 		self.test_collection = self._collections[Collections.SESSION]
 		session_repo.init(collection_name=self.test_collection, drop=True)
-
 		self.account_id = str(ObjectId())
 		self.patient_id = str(ObjectId())
 
@@ -30,155 +32,146 @@ class TestSessionRepository(BaseMongoTest):
 
 	def test_create_and_get_session(self):
 		"""Test chat session creation and retrieval by ID."""
-		# Test creation
 		session = session_repo.create_session(
-			self.account_id,
-			self.patient_id,
-			"Test Chat",
-			collection_name=self.test_collection
+			self.account_id, self.patient_id, "Test Chat", collection_name=self.test_collection
 		)
-		self.assertIn("_id", session)
 		self.assertIsInstance(session["_id"], str)
-		self.assertEqual(session["title"], "Test Chat")
-		self.assertEqual(len(session["messages"]), 0)
-
-		# Test retrieval
 		retrieved = session_repo.get_session(session["_id"], collection_name=self.test_collection)
 		self.assertIsNotNone(retrieved)
 		self.assertEqual(retrieved["_id"], session["_id"]) # type: ignore
-		self.assertEqual(retrieved["account_id"], self.account_id) # type: ignore
-		self.assertEqual(retrieved["patient_id"], self.patient_id) # type: ignore
-
 		# Test getting a non-existent session
-		non_existent = session_repo.get_session(str(ObjectId()), collection_name=self.test_collection)
-		self.assertIsNone(non_existent)
+		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."""
-		session = session_repo.create_session(
-			self.account_id, self.patient_id, "Message Test", collection_name=self.test_collection
-		)
-		session_id = session["_id"]
-
-		# Add messages and verify session's updated_at timestamp changes
-		original_doc = self.get_doc_by_id(Collections.SESSION, session_id)
-		time.sleep(0.01) # Ensure timestamp will be different
-		session_repo.add_message(session_id, "User message 1", True, collection_name=self.test_collection)
-		updated_doc = self.get_doc_by_id(Collections.SESSION, session_id)
-		self.assertLess(original_doc["updated_at"], updated_doc["updated_at"]) # type: ignore
-
-		session_repo.add_message(session_id, "AI response 1", False, collection_name=self.test_collection)
-		session_repo.add_message(session_id, "User message 2", True, collection_name=self.test_collection)
-
-		# Test message retrieval (should be in descending order of creation)
-		messages = session_repo.get_session_messages(session_id, collection_name=self.test_collection)
-		self.assertEqual(len(messages), 3)
-		self.assertEqual(messages[0]["_id"], 2)
-		self.assertEqual(messages[0]["content"], "User message 2")
-		self.assertEqual(messages[1]["_id"], 1)
-		self.assertEqual(messages[2]["_id"], 0)
-
+		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)
+		self.assertEqual(len(messages), 2)
+		self.assertEqual(messages[0]["content"], "AI response 1") # Descending order
+		self.assertEqual(messages[1]["_id"], 0)
 		# Test limit
-		limited_messages = session_repo.get_session_messages(session_id, limit=2, collection_name=self.test_collection)
-		self.assertEqual(len(limited_messages), 2)
-		self.assertEqual(limited_messages[0]["_id"], 2)
-
-		# Test adding message to non-existent session
-		with self.assertRaises(ActionFailed):
-			session_repo.add_message(str(ObjectId()), "ghost", True, collection_name=self.test_collection)
-
-	def test_list_patient_sessions(self):
-		"""Test listing sessions for a specific patient, sorted by update time."""
-		p_id_1 = str(ObjectId())
-		p_id_2 = str(ObjectId())
-
-		# Create sessions, sleeping briefly to ensure distinct updated_at times
-		session_repo.create_session(self.account_id, p_id_1, "P1 Chat 1", collection_name=self.test_collection)
-		time.sleep(0.01)
-		session_repo.create_session(self.account_id, p_id_2, "P2 Chat 1", collection_name=self.test_collection) # Belongs to other patient
-		time.sleep(0.01)
-		s2 = session_repo.create_session(self.account_id, p_id_1, "P1 Chat 2", collection_name=self.test_collection)
-
-		# Test listing for patient 1
-		sessions = session_repo.list_patient_sessions(p_id_1, collection_name=self.test_collection)
-		self.assertEqual(len(sessions), 2)
-		self.assertEqual(sessions[0]["_id"], s2["_id"]) # Most recently created should be first
-
-	def test_get_user_sessions(self):
-		"""Test listing sessions for a specific user, sorted by update time."""
-		user1 = str(ObjectId())
-		user2 = str(ObjectId())
-
-		s1 = session_repo.create_session(user1, self.patient_id, "U1 Chat 1", collection_name=self.test_collection)
-		time.sleep(0.01)
-		session_repo.create_session(user2, self.patient_id, "U2 Chat 1", collection_name=self.test_collection)
-		time.sleep(0.01)
-		s3 = session_repo.create_session(user1, self.patient_id, "U1 Chat 2", collection_name=self.test_collection)
-
-		sessions = session_repo.get_user_sessions(user1, collection_name=self.test_collection)
-		self.assertEqual(len(sessions), 2)
-		self.assertEqual(sessions[0]["_id"], s3["_id"]) # s3 was updated most recently
-		self.assertEqual(sessions[1]["_id"], s1["_id"])
-
-		# Test limit
-		sessions_limited = session_repo.get_user_sessions(user1, limit=1, collection_name=self.test_collection)
-		self.assertEqual(len(sessions_limited), 1)
+		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."""
+		session_repo.create_session(self.account_id, self.patient_id, "First", collection_name=self.test_collection)
+		time.sleep(0.01) # Ensure distinct timestamps
+		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
+		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"])
 
 	def test_update_session_title(self):
-		"""Test updating a session's title."""
-		session = session_repo.create_session(self.account_id, self.patient_id, "Old Title", collection_name=self.test_collection)
-		session_id = session["_id"]
-
-		success = session_repo.update_session_title(session_id, "New Title", collection_name=self.test_collection)
+		"""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)
 		self.assertTrue(success)
-
-		updated_session = session_repo.get_session(session_id, collection_name=self.test_collection)
-		self.assertEqual(updated_session["title"], "New Title") # type: ignore
-
-		# Test updating non-existent session
-		success_fail = session_repo.update_session_title(str(ObjectId()), "Ghost", collection_name=self.test_collection)
-		self.assertFalse(success_fail)
+		updated_doc = self.get_doc_by_id(Collections.SESSION, session["_id"])
+		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))
 
 	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)
-		session_id = session["_id"]
-
-		success = session_repo.delete_session(session_id, collection_name=self.test_collection)
-		self.assertTrue(success)
-
-		deleted_session = session_repo.get_session(session_id, collection_name=self.test_collection)
-		self.assertIsNone(deleted_session)
-
-		# Test deleting non-existent session
-		success_fail = session_repo.delete_session(str(ObjectId()), collection_name=self.test_collection)
-		self.assertFalse(success_fail)
+		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."""
-		coll = get_collection(self.test_collection)
-		now = datetime.now(timezone.utc)
-		old_date = now - timedelta(days=31)
-
-		# Manually insert one old and one new session
-		coll.insert_one({
-			"account_id": ObjectId(self.account_id), "patient_id": ObjectId(self.patient_id),
-			"title": "Old Session", "created_at": old_date, "updated_at": old_date, "messages": []
-		})
-		coll.insert_one({
-			"account_id": ObjectId(self.account_id), "patient_id": ObjectId(self.patient_id),
-			"title": "New Session", "created_at": now, "updated_at": now, "messages": []
-		})
-
-		self.assertEqual(coll.count_documents({}), 2)
+		# 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_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}}
+		)
 
+		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(coll.count_documents({}), 1)
+		self.assertEqual(get_collection(self.test_collection).count_documents({}), 1)
 
-		remaining = coll.find_one()
-		self.assertEqual(remaining["title"], "New Session") # type: ignore
 
+class TestSessionRepositoryExceptions(BaseMongoTest):
+	"""Test class for the exception handling of all session repository functions."""
+
+	def setUp(self):
+		"""Set up a clean test environment before each test."""
+		super().setUp()
+		self.test_collection = self._collections[Collections.SESSION]
+		session_repo.init(collection_name=self.test_collection, drop=True)
+		self.account_id = str(ObjectId())
+		self.patient_id = str(ObjectId())
+
+	@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
+
+		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)
+
+	def test_invalid_id_raises_action_failed(self):
+		"""Test that functions raise ActionFailed when given a malformed ObjectId string."""
+		with self.assertRaises(ActionFailed):
+			session_repo.create_session("bad-id", self.patient_id, "t", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.get_user_sessions("bad-id", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.list_patient_sessions("bad-id", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.get_session("bad-id", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.get_session_messages("bad-id", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.update_session_title("bad-id", "t", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.delete_session("bad-id", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.add_message("bad-id", "t", True, collection_name=self.test_collection)
+
+	@patch('src.data.repositories.session.get_collection')
+	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):
+			session_repo.create_session(self.account_id, self.patient_id, "t", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.get_user_sessions(self.account_id, collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.list_patient_sessions(self.patient_id, collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.get_session(str(ObjectId()), collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.get_session_messages(str(ObjectId()), collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.update_session_title(str(ObjectId()), "t", collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.delete_session(str(ObjectId()), collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.prune_old_sessions(collection_name=self.test_collection)
+		with self.assertRaises(ActionFailed):
+			session_repo.add_message(str(ObjectId()), "t", True, collection_name=self.test_collection)
 
 if __name__ == "__main__":
 	try: