tests_updated_project_complete

Author: RobuRishabh

README.md

@@ -4,14 +4,27 @@
 The RAG-based Flight Query Bot is a Streamlit-powered web application that allows users to query flight information using natural language (e.g., "Show me flights from New York to London"). It uses Retrieval-Augmented Generation (RAG) with an Ollama language model to extract entities from queries and retrieve flight data from a mock database. The project is deployed on a local Kubernetes cluster using Minikube and includes a CI/CD pipeline with GitHub Actions for automated testing.
 
 ### Architecture
+```
+  RAG-based-Flight-Query-Bot/
+  ├── .github/
+  │   └── workflows/
+  │       └── test.yml
+  ├── tests/
+  │   └── test_mock_database.py
+  ├── mock_database.py
+  ├── query_handler.py
+  ├── ollama_api.py
+  ├── app.py (optional)
+  ├── requirements.txt
+  └── README.md
+  ```
 - **Frontend**: Streamlit UI (`app.py`) for user interaction.
 - **Backend**: 
   - `query_handler.py`: Processes queries and extracts entities using Ollama.
   - `ollama_api.py`: Integrates with the Ollama LLM for natural language responses.
   - `mock_database.py`: Provides mock flight data and search functionality.
 - **Deployment**: Kubernetes on Minikube with two services: `flight-assistant-service` (Streamlit) and `ollama-service` (Ollama server).
 - **CI/CD**: GitHub Actions runs unit tests on every push or pull request.
-
 ---
 
 ## Setup Instructions

tests/test_mock_databse.py

@@ -1,36 +1,46 @@
 import unittest
+from unittest.mock import patch
+import os
+import requests
 from mock_database import flights, search_flights, check_ollama_availability
 
 class TestMockDatabase(unittest.TestCase):
     def test_flight_data_exists(self):
         self.assertGreater(len(flights), 0, "Flight database should not be empty")
 
     def test_flight_structure(self):
-        flight = flights[0]
+        self.assertEqual(len(flights), 5, "Expected exactly 5 flights")
         expected_keys = {"flight_number", "origin", "destination", "time", "airline"}
-        self.assertEqual(set(flight.keys()), expected_keys, "Flight data has incorrect structure")
+        for flight in flights:
+            self.assertEqual(set(flight.keys()), expected_keys, "Flight data has incorrect structure")
+            self.assertIsInstance(flight["flight_number"], str, "Flight number should be a string")
+            self.assertNotEqual(flight["origin"], flight["destination"], "Origin and destination should differ")
 
     def test_search_flights_by_flight_number(self):
-        results = search_flights(flight_number="NY100")  # Use keyword arg for flight_number
-        self.assertGreater(len(results), 0, "Should find at least one matching flight")
-        self.assertIn("NY100", [f["flight_number"] for f in results], "Flight NY100 should be found")
+        results = search_flights(flight_number="NY100")
         self.assertEqual(len(results), 1, "Should find exactly one flight for NY100")
+        self.assertEqual(results[0]["flight_number"], "NY100", "Flight NY100 should be found")
 
     def test_search_flights_by_origin(self):
         results = search_flights(origin="New York")
-        self.assertGreater(len(results), 0, "Should find flights from New York")
+        self.assertEqual(len(results), 1, "Should find exactly one flight from New York")
         self.assertTrue(all(f["origin"] == "New York" for f in results), "All results should have origin New York")
 
     def test_search_flights_by_destination(self):
         results = search_flights(destination="London")
-        self.assertGreater(len(results), 0, "Should find flights to London")
+        self.assertEqual(len(results), 1, "Should find exactly one flight to London")
         self.assertTrue(all(f["destination"] == "London" for f in results), "All results should have destination London")
 
     def test_search_flights_by_airline(self):
         results = search_flights(airline="Global Airways")
-        self.assertGreater(len(results), 0, "Should find flights by Global Airways")
+        self.assertEqual(len(results), 1, "Should find exactly one flight by Global Airways")
         self.assertTrue(all(f["airline"] == "Global Airways" for f in results), "All results should have airline Global Airways")
 
+    def test_search_flights_combined_filters(self):
+        results = search_flights(origin="San Francisco", destination="Sydney")
+        self.assertEqual(len(results), 1, "Should find exactly one flight from San Francisco to Sydney")
+        self.assertEqual(results[0]["flight_number"], "SF400", "Flight SF400 should be found")
+
     def test_search_flights_no_results(self):
         results = search_flights(flight_number="XYZ999")
         self.assertEqual(len(results), 0, "Should return no results for invalid flight number")
@@ -41,8 +51,29 @@ def test_search_flights_no_parameters(self):
 
     def test_search_flights_case_insensitivity(self):
         results = search_flights(flight_number="ny100")
-        self.assertGreater(len(results), 0, "Should find NY100 regardless of case")
-        self.assertIn("NY100", [f["flight_number"] for f in results], "Flight NY100 should be found")
+        self.assertEqual(len(results), 1, "Should find NY100 regardless of case")
+        self.assertEqual(results[0]["flight_number"], "NY100", "Flight NY100 should be found")
+
+    def test_search_flights_ignore_city_name(self):
+        results = search_flights(destination="City Name")
+        self.assertEqual(len(results), 0, "Should ignore 'City Name' as destination")
+
+    @patch("requests.get")
+    def test_check_ollama_availability_success(self, mock_get):
+        mock_response = mock_get.return_value
+        mock_response.status_code = 200
+        self.assertTrue(check_ollama_availability(), "Should return True when Ollama is available")
+
+    @patch("requests.get")
+    def test_check_ollama_availability_failure(self, mock_get):
+        mock_response = mock_get.return_value
+        mock_response.status_code = 500
+        self.assertFalse(check_ollama_availability(), "Should return False on server error")
+
+    @patch("requests.get")
+    def test_check_ollama_availability_exception(self, mock_get):
+        mock_get.side_effect = requests.RequestException("Connection error")
+        self.assertFalse(check_ollama_availability(), "Should return False on exception")
 
 if __name__ == "__main__":
     unittest.main()

tests/test_ollama_api.py

@@ -1,94 +1,109 @@
-import unittest
+import pytest
 import requests
 from unittest.mock import patch, Mock
-from ollama_api import initialize_ollama, check_ollama_availability, generate_response, generate_fallback_response
+import os
+from ollama_api import initialize_ollama, check_ollama_availability, generate_fallback_response, generate_response, ollama_llm
 
-class TestOllamaAPI(unittest.TestCase):
-    def setUp(self):
-        # Reset global ollama_llm for each test
-        global ollama_llm
-        ollama_llm = None
+# Fixture to mock environment variables
+@pytest.fixture
+def mock_env():
+    original_env = os.environ.copy()
+    os.environ["OLLAMA_URL"] = "http://test:11434"
+    os.environ["OLLAMA_MODEL"] = "llama2:test"
+    yield
+    os.environ.clear()
+    os.environ.update(original_env)
 
-    @patch("ollama_api.OllamaLLM")
-    def test_initialize_ollama_success(self, mock_ollama_llm):
-        mock_llm_instance = Mock()
-        mock_ollama_llm.return_value = mock_llm_instance
-        result = initialize_ollama()
-        self.assertIsNotNone(result, "Ollama LLM should initialize successfully")
-        self.assertEqual(result, mock_llm_instance)
+# 1. Tests for initialize_ollama
+@patch("ollama_api.OllamaLLM")
+def test_initialize_ollama_success(mock_ollama, mock_env):
+    mock_instance = Mock()
+    mock_ollama.return_value = mock_instance
+    result = initialize_ollama()
+    assert result == mock_instance, "Should return OllamaLLM instance on success"
+    mock_ollama.assert_called_once_with(model="llama2:test", base_url="http://test:11434")
 
-    @patch("ollama_api.OllamaLLM", side_effect=Exception("Initialization failed"))
-    def test_initialize_ollama_failure(self, mock_ollama_llm):
-        result = initialize_ollama()
-        self.assertIsNone(result, "Ollama LLM should return None on failure")
+@patch("ollama_api.OllamaLLM")
+def test_initialize_ollama_failure(mock_ollama, mock_env):
+    mock_ollama.side_effect = Exception("Connection failed")
+    result = initialize_ollama()
+    assert result is None, "Should return None on initialization failure"
 
-    @patch("ollama_api.requests.get")
-    def test_check_ollama_availability_success(self, mock_get):
-        mock_response = Mock()
-        mock_response.status_code = 200
-        mock_get.return_value = mock_response
-        is_available, message = check_ollama_availability()
-        self.assertTrue(is_available, "Ollama should be available")
-        self.assertIn("available", message)
+# 2. Tests for check_ollama_availability
+@patch("requests.get")
+def test_check_ollama_availability_success(mock_get, mock_env):
+    mock_response = mock_get.return_value
+    mock_response.status_code = 200
+    is_available, message = check_ollama_availability()
+    assert is_available is True, "Should return True when server is available"
+    assert "Ollama server is available" in message, "Should return success message"
+    mock_get.assert_called_once_with("http://test:11434/api/tags", timeout=3)
 
-    @patch("ollama_api.requests.get", side_effect=requests.RequestException("Connection error"))
-    def test_check_ollama_availability_failure(self, mock_get):
-        is_available, message = check_ollama_availability()
-        self.assertFalse(is_available, "Ollama should not be available")
-        self.assertIn("not available", message)
+@patch("requests.get")
+def test_check_ollama_availability_failure(mock_get, mock_env):
+    mock_get.side_effect = requests.RequestException("Connection error")
+    is_available, message = check_ollama_availability()
+    assert is_available is False, "Should return False on request exception"
+    assert "Ollama server is not available" in message, "Should return error message"
 
-    def test_generate_fallback_response_with_flights(self):
-        flights = [
-            {"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"}
-        ]
-        response = generate_fallback_response("test query", flights)
-        self.assertIn("NY100", response)
-        self.assertIn("New York", response)
-        self.assertIn("London", response)
-        self.assertIn("2025-05-01 08:00", response)
-        self.assertIn("Global Airways", response)
+# 3. Tests for generate_fallback_response
+def test_generate_fallback_response_with_flights():
+    flights = [
+        {"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"}
+    ]
+    result = generate_fallback_response("flights from New York", flights)
+    assert "Flight NY100" in result, "Should include flight number"
+    assert "New York to London" in result, "Should include origin and destination"
+    assert "Time: 2025-05-01 08:00" in result, "Should include time"
+    assert "Airline: Global Airways" in result, "Should include airline"
 
-    def test_generate_fallback_response_no_flights(self):
-        response = generate_fallback_response("test query", [])
-        self.assertEqual(response, "I couldn't find any flights matching your criteria. Please try again.")
+def test_generate_fallback_response_no_flights():
+    result = generate_fallback_response("flights from Mars", [])
+    assert "I couldn't find any flights" in result, "Should return no-flights message"
 
-    @patch("ollama_api.check_ollama_availability")
-    @patch("ollama_api.ollama_llm", new=Mock(invoke=lambda x: "Flight NY100 departs at 08:00 from New York to London with Global Airways"))
-    def test_generate_response_success(self, mock_check_availability):
-        mock_check_availability.return_value = (True, "Ollama available")
-        query = "Show me flight NY100"
-        flight_data = [
-            {"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"}
-        ]
-        response = generate_response(query, flight_data)
-        self.assertIn("NY100", response)
-        self.assertIn("08:00", response)
-        self.assertIn("New York", response)
-        self.assertIn("London", response)
-        self.assertIn("Global Airways", response)
+def test_generate_fallback_response_missing_fields():
+    flights = [{"flight_number": "NY100"}]  # Missing other fields
+    result = generate_fallback_response("test query", flights)
+    assert "Flight NY100" in result, "Should include flight number"
+    assert "Unknown to Unknown" in result, "Should handle missing origin/destination"
+    assert "Time: N/A" in result, "Should handle missing time"
+    assert "Airline: N/A" in result, "Should handle missing airline"
 
-    @patch("ollama_api.check_ollama_availability")
-    def test_generate_response_no_flights(self, mock_check_availability):
-        mock_check_availability.return_value = (True, "Ollama available")
-        # Ensure ollama_llm is initialized (mock it globally)
-        global ollama_llm
-        ollama_llm = Mock(invoke=lambda x: "No flights found.")
-        query = "Show me flight XYZ999"
-        flight_data = []
-        response = generate_response(query, flight_data)
-        self.assertEqual(response, "I couldn't find any flights matching your criteria. Please try again.")
+# 4. Tests for generate_response
+@patch("ollama_api.check_ollama_availability")
+@patch("ollama_api.ollama_llm.invoke")
+def test_generate_response_ollama_success(mock_invoke, mock_check, mock_env):
+    mock_check.return_value = (True, "Server available")
+    mock_invoke.return_value = "Flight NY100 departs from New York to London at 08:00 with Global Airways."
+    with patch("ollama_api.ollama_llm", Mock()):
+        flights = [{"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"}]
+        result = generate_response("flights from New York", flights)
+        assert "NY100" in result, "Should include flight details from Ollama"
+        assert "New York to London" in result, "Should include route"
+        mock_invoke.assert_called_once()
 
-    @patch("ollama_api.check_ollama_availability")
-    def test_generate_response_ollama_unavailable(self, mock_check_availability):
-        mock_check_availability.return_value = (False, "Ollama not available")
-        query = "Show me flight NY100"
-        flight_data = [
-            {"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"}
-        ]
-        response = generate_response(query, flight_data)
-        self.assertIn("NY100", response)
-        self.assertIn("New York", response)
-        self.assertIn("London", response)
+@patch("ollama_api.check_ollama_availability")
+def test_generate_response_ollama_unavailable(mock_check, mock_env):
+    mock_check.return_value = (False, "Server unavailable")
+    flights = [{"flight_number": "NY100", "origin": "New York", "destination": "London"}]
+    result = generate_response("flights from New York", flights)
+    assert "Flight NY100" in result, "Should use fallback when Ollama unavailable"
+    assert "New York to London" in result, "Should include route in fallback"
 
-if __name__ == "__main__":
-    unittest.main()
+@patch("ollama_api.check_ollama_availability")
+def test_generate_response_ollama_not_initialized(mock_check, mock_env):
+    mock_check.return_value = (True, "Server available")
+    with patch("ollama_api.ollama_llm", None):
+        flights = [{"flight_number": "NY100"}]
+        result = generate_response("test query", flights)
+        assert "Flight NY100" in result, "Should use fallback when ollama_llm is None"
+
+@patch("ollama_api.check_ollama_availability")
+@patch("ollama_api.ollama_llm.invoke")
+def test_generate_response_ollama_failure(mock_invoke, mock_check, mock_env):
+    mock_check.return_value = (True, "Server available")
+    mock_invoke.side_effect = Exception("LLM error")
+    with patch("ollama_api.ollama_llm", Mock()):
+        flights = [{"flight_number": "NY100"}]
+        result = generate_response("test query", flights)
+        assert "Flight NY100" in result, "Should use fallback on Ollama exception"