diff --git a/agentic-rag-authorization/ARCHITECTURE.md b/agentic-rag-authorization/ARCHITECTURE.md
index 58ac9ce..2fb0094 100644
--- a/agentic-rag-authorization/ARCHITECTURE.md
+++ b/agentic-rag-authorization/ARCHITECTURE.md
@@ -1258,7 +1258,7 @@ agentic-rag-elasticsearch/
 │
 ├── run_ui.py                  # Web UI launcher with pre-flight checks
 ├── docker-compose.yml         # Elasticsearch + SpiceDB services
-├── requirements.txt           # Python dependencies (includes fastapi, uvicorn)
+├── pyproject.toml             # Project metadata + dependencies ([ui] extra adds fastapi, uvicorn)
 ├── .env.example               # Environment variable template
 ├── README.md                  # Overview, quick start, usage guide
 └── ARCHITECTURE.md            # This file: deep technical dive
diff --git a/agentic-rag-authorization/README.md b/agentic-rag-authorization/README.md
index 0e2156d..50b49ed 100644
--- a/agentic-rag-authorization/README.md
+++ b/agentic-rag-authorization/README.md
@@ -1,6 +1,6 @@
 # Agentic RAG with Fine-Grained Authorization
 
-This repository demonstrates how to combine agentic behavior with deterministic fine-grained authorization using LangGraph, SpiceDB, [Elasticsearch](https://github.com/elastic/elasticsearch), and [Mistral](https://mistral.ai/). You'll learn to build RAG systems where a user can only see information from the documents they have access to.
+This repository demonstrates how to combine agentic behavior with deterministic fine-grained authorization using LangGraph, SpiceDB, [Elasticsearch](https://github.com/elastic/elasticsearch), and [Mistral](https://mistral.ai/). You'll learn to build an Agentic RAG system where a user can only see information from the documents they have access to.
 
 This project uses the [LangChain SpiceDB](https://pypi.org/project/langchain-spicedb/) library.
 
@@ -9,13 +9,13 @@ This project uses the [LangChain SpiceDB](https://pypi.org/project/langchain-spi
 
 ## TL;DR (human-written)
 
-RAG systems typically focus on the retrieval mechanisms, but don't have fine-grained access control to check if the information retrieved is accessible to the user asking the query. This demo shows the setup for a prod-like Agentic RAG. It has a corpus of 50 documents with complex sharing requirements that span individual, departments and exceptions. 
+RAG systems typically focus on the retrieval mechanisms, but don't have fine-grained access control to check if the information retrieved is accessible to the user asking the query. This demo shows the setup for a prod-like Agentic RAG system. It has a corpus of 50 documents with [complex sharing requirements](https://github.com/authzed/examples/blob/main/agentic-rag-authorization/data/PERMISSIONS.md) that span individual, departments and exceptions. 
 
 The two takeaways from this demo are:
 
-1. Using ReBAC makes it simple to model complex hierarchal permissions. The complexity increases in the context of RAG and AI Applications as there are 10x more principals, so traditional authorization methods such as RBAC fall flat.
+1. Using ReBAC makes it simple to model complex hierarchal permissions. The complexity increases in the context of RAG and AI Applications as there are 10x more principals, so traditional authorization methods such as RBAC, row-level security, or scoped tokens fall flat.
 
-2. Never ever let an AI Agent *decide* if it needs to check for authorization. Gen AI is inherently probabilistic so you have to ensure that permission checks are deterministic and cannot be skipped.
+2. Never ever let an AI Agent *decide* if it needs to check for authorization. Generative AI is inherently probabilistic so you have to ensure that permission checks are deterministic and cannot be skipped or else you might fall victim to prompt injection attacks.
 
 The diagram below shows the end-to-end request flow: the application embeds the query with Mistral, retrieves relevant context from Elasticsearch, checks permissions in SpiceDB, and only then passes the authorized context to the Mistral LLM to generate a response.
 
@@ -25,19 +25,20 @@ The diagram below shows the end-to-end request flow: the application embeds the
 
 This repo demonstrates:
 
-1. **Fine-grained authorization in Agentic RAG** - How to enforce document-level permissions with SpiceDB so users only see what they're allowed to see
+1. **Fine-grained authorization in Agentic RAG** - How to enforce document-level permissions with SpiceDB so users get responses based only on what they're allowed to see
 2. **Security architecture** - A deterministic authorization boundary that cannot be bypassed by the agent
 3. **Production features** - Structured logging, connection pooling, batch operations, error handling
 4. **Real-world complexity** - 50 documents, 4 permission patterns with hierarchies
 
-Note: Despite the "agentic RAG" name, the default mode is intentionally simple and deterministic (3 nodes: retrieve → authorize → generate). This provides fast, predictable behavior suitable for most use cases. There is a `MAX_RETRIEVAL_ATTEMPTS` option where the AI Agent can reason if it has to retrieve more data.
+Note: Despite the "agentic RAG" name, the default mode is intentionally simple and deterministic (3 nodes: retrieve → authorize → generate). This provides fast, predictable behavior suitable for most use cases. There is a `MAX_RETRIEVAL_ATTEMPTS` option where the AI Agent can **reason** if it has to retrieve more data.
 
 ## The Problem This Solves
 
 Traditional RAG retrieves documents by semantic similarity without considering permissions. This creates two issues:
 
-1. **Security risk**: Users might see documents they shouldn't access
-2. **Poor UX**: Silent failures when documents are denied, with no explanation
+1. **Security risk**: Users might see documents they shouldn't access leading to data leakage. This can lead to reputation and customer damage if data is breached through injection attacks.
+2. **Audit and regulatory concerns** - Building enterprise-ready AI requires audit logs especially in compliance-regulated industries. 
+3. **Poor UX**: Silent failures when documents are denied, with no explanation
 
 Read the [OWASP Top 10 for LLM](https://owasp.org/www-project-top-10-for-large-language-model-applications/) and [OWASP Top 10 Risks to Web Apps](https://owasp.org/Top10/2025/A01_2025-Broken_Access_Control/) for more on why access control matters.
 
@@ -93,9 +94,9 @@ The agent transparently explains access limitations instead of failing silently.
 The demo runs entirely through the web UI, which lets you switch users and watch the authorization boundary filter results in real time.
 
 ### Prerequisites
-- Docker & Docker Compose
+- Docker (runs SpiceDB and Elasticsearch containers)
 - Python 3.11+
-- Mistral API key
+- Mistral (or any other LLM) API key
 
 ### Steps
 
@@ -106,10 +107,10 @@ cp .env.example .env
 # 2. Start Elasticsearch + SpiceDB
 docker-compose up -d
 
-# 3. Install dependencies (includes FastAPI + Uvicorn for the UI)
+# 3. Install dependencies (the [ui] extra adds FastAPI + Uvicorn for the web UI)
 python3 -m venv venv
 source venv/bin/activate
-pip install -r requirements.txt
+pip install -e ".[ui]"
 
 # 4. Initialize data — embed 50 documents into Elasticsearch
 #    and write the permission model to SpiceDB
@@ -123,9 +124,19 @@ python3 run_ui.py
 
 `run_ui.py` verifies Elasticsearch and SpiceDB connectivity, confirms the documents are loaded, starts the FastAPI server, and opens http://localhost:8000.
 
+### Running the tests
+
+```bash
+pip install -e ".[dev]"
+pytest                 # unit tests — no services required
+pytest -m integration  # live tests against Elasticsearch + SpiceDB
+```
+
 ## How It Works
 
-### 1. Authorization Model (SpiceDB)
+### 1. Authorization Model 
+
+#### Here's the SpiceDB Schema:
 
 ```zed
 definition user {}
@@ -196,35 +207,6 @@ MAX_RETRIEVAL_ATTEMPTS=3  # default is 1
 - **Fail closed**: Access denied unless explicitly granted
 - **Observable**: Full audit trail in state
 
-## Configuration
-
-Environment variables (`.env`):
-
-```bash
-# Required
-MISTRAL_API_KEY=...
-
-# Optional (defaults shown)
-ELASTICSEARCH_URL=http://localhost:9200
-ELASTICSEARCH_API_KEY=
-SPICEDB_ENDPOINT=localhost:50051
-SPICEDB_TOKEN=devtoken
-MAX_RETRIEVAL_ATTEMPTS=1
-LOG_LEVEL=INFO
-```
-
-## Dataset Overview
-
-The repository includes a realistic 50-document dataset across 5 departments.
-
-**Authorization Patterns:**
-1. Department-based access (primary pattern)
-2. Cross-department collaboration (3 shared documents)
-3. Individual user exceptions (3 special grants)
-4. Public documents (accessible to all users)
-
-See [data/PERMISSIONS.md](data/PERMISSIONS.md) for the complete permission matrix.
-
 ## Learn More
 
 - **SpiceDB**: https://authzed.com/docs
diff --git a/agentic-rag-authorization/docker-compose.yml b/agentic-rag-authorization/docker-compose.yml
index f175890..429623a 100644
--- a/agentic-rag-authorization/docker-compose.yml
+++ b/agentic-rag-authorization/docker-compose.yml
@@ -2,7 +2,7 @@
 services:
   elasticsearch:
     container_name: "elasticsearch"
-    image: "docker.elastic.co/elasticsearch/elasticsearch:8.13.4"
+    image: "docker.elastic.co/elasticsearch/elasticsearch:9.5.2"
     environment:
       - "discovery.type=single-node"
       - "xpack.security.enabled=false"
diff --git a/agentic-rag-authorization/pyproject.toml b/agentic-rag-authorization/pyproject.toml
index 4c8c8a4..cea27d7 100644
--- a/agentic-rag-authorization/pyproject.toml
+++ b/agentic-rag-authorization/pyproject.toml
@@ -4,28 +4,38 @@ version = "0.1.0"
 description = "Agentic RAG with fine-grained authorization using Elasticsearch and SpiceDB"
 requires-python = ">=3.10"
 dependencies = [
-    "langchain>=0.1.0",
     "langchain-mistralai>=0.2.0",
     "langgraph>=0.0.20",
-    "langchain-spicedb>=0.1.0",
-    "elasticsearch>=8.0.0,<9.0.0",
+    "langchain-spicedb>=0.2.0",
+    "elasticsearch>=9.0.0,<10.0.0",
     "authzed>=0.7.0",
     "python-dotenv>=1.0.0",
 ]
 
 [project.optional-dependencies]
+ui = [
+    "fastapi>=0.109.0",
+    "uvicorn[standard]>=0.27.0",
+]
 dev = [
     "pytest>=7.0.0",
     "pytest-asyncio>=0.21.0",
     "ruff>=0.1.0",
 ]
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "strict"
+markers = [
+    "integration: tests requiring live Elasticsearch/SpiceDB/Mistral services (deselect with '-m \"not integration\"')",
+]
+
 [build-system]
 requires = ["setuptools>=61.0", "wheel"]
-build-backend = "setuptools.build_backend"
+build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
-packages = ["agentic_rag", "agentic_rag.nodes", "agentic_rag.tools"]
+packages = ["agentic_rag", "agentic_rag.nodes"]
 
 [tool.setuptools.package-data]
 agentic_rag = ["py.typed"]
diff --git a/agentic-rag-authorization/requirements.txt b/agentic-rag-authorization/requirements.txt
deleted file mode 100644
index 767fa56..0000000
--- a/agentic-rag-authorization/requirements.txt
+++ /dev/null
@@ -1,19 +0,0 @@
-# Core dependencies
-langchain>=0.1.0
-langchain-mistralai>=0.2.0
-langgraph>=0.0.20
-langchain-spicedb>=0.2.0
-elasticsearch>=8.0.0,<9.0.0
-authzed>=0.7.0
-python-dotenv>=1.0.0
-grpcio>=1.50.0
-grpcio-tools>=1.50.0
-
-# Web UI dependencies
-fastapi>=0.109.0
-uvicorn[standard]>=0.27.0
-
-# Development dependencies
-pytest>=7.0.0
-pytest-asyncio>=0.21.0
-ruff>=0.1.0
diff --git a/agentic-rag-authorization/test_dataset_integration.py b/agentic-rag-authorization/test_dataset_integration.py
deleted file mode 100644
index 6c7641d..0000000
--- a/agentic-rag-authorization/test_dataset_integration.py
+++ /dev/null
@@ -1,171 +0,0 @@
-"""Integration test for the realistic document dataset implementation."""
-
-import sys
-import os
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'scripts'))
-
-from parse_documents import load_all_documents
-
-
-def test_document_loading():
-    """Test that documents can be loaded correctly."""
-    print("Testing document loading...")
-
-    docs = load_all_documents()
-
-    # Test count
-    assert len(docs) == 50, f"Expected 50 documents, got {len(docs)}"
-    print(f"  ✅ Loaded {len(docs)} documents")
-
-    # Test required fields
-    for doc in docs:
-        assert 'doc_id' in doc, f"Missing doc_id in {doc}"
-        assert 'title' in doc, f"Missing title in {doc}"
-        assert 'content' in doc, f"Missing content in {doc}"
-        assert 'department' in doc, f"Missing department in {doc}"
-        assert 'classification' in doc, f"Missing classification in {doc}"
-
-    print("  ✅ All documents have required fields")
-
-    return docs
-
-
-def test_department_distribution(docs):
-    """Test that document distribution matches expectations."""
-    print("\nTesting department distribution...")
-
-    dept_counts = {}
-    for doc in docs:
-        dept = doc['department']
-        dept_counts[dept] = dept_counts.get(dept, 0) + 1
-
-    expected = {
-        'engineering': 15,
-        'sales': 10,
-        'hr': 10,
-        'finance': 10,
-        'public': 5,
-    }
-
-    for dept, expected_count in expected.items():
-        actual_count = dept_counts.get(dept, 0)
-        assert actual_count == expected_count, \
-            f"{dept}: expected {expected_count}, got {actual_count}"
-        print(f"  ✅ {dept}: {actual_count} documents")
-
-
-def test_document_structure(docs):
-    """Test document content structure."""
-    print("\nTesting document structure...")
-
-    # Check engineering document
-    eng_docs = [d for d in docs if d['department'] == 'engineering']
-    assert len(eng_docs) > 0, "No engineering documents found"
-
-    sample_eng = eng_docs[0]
-    assert 'Architecture' in sample_eng['title'] or 'Guide' in sample_eng['title'] \
-        or 'Memo' in sample_eng['title'] or 'Specification' in sample_eng['title'], \
-        f"Unexpected title format: {sample_eng['title']}"
-    print(f"  ✅ Engineering doc sample: {sample_eng['title']}")
-
-    # Check sales document
-    sales_docs = [d for d in docs if d['department'] == 'sales']
-    assert len(sales_docs) > 0, "No sales documents found"
-
-    sample_sales = sales_docs[0]
-    print(f"  ✅ Sales doc sample: {sample_sales['title']}")
-
-    # Check public document
-    public_docs = [d for d in docs if d['department'] == 'public']
-    assert len(public_docs) == 5, f"Expected 5 public docs, got {len(public_docs)}"
-
-    for pub_doc in public_docs:
-        assert pub_doc['classification'] == 'public', \
-            f"Public doc has wrong classification: {pub_doc['classification']}"
-
-    print(f"  ✅ All {len(public_docs)} public documents have 'public' classification")
-
-
-def test_cross_department_docs(docs):
-    """Test that cross-department documents exist."""
-    print("\nTesting cross-department documents...")
-
-    # These are the cross-department docs we configured
-    cross_dept_ids = [
-        'engineering-architecture-001',
-        'sales-guide-005',
-        'hr-policy-001',
-    ]
-
-    for doc_id in cross_dept_ids:
-        found = any(d['doc_id'] == doc_id for d in docs)
-        assert found, f"Cross-department document not found: {doc_id}"
-        print(f"  ✅ Found cross-dept doc: {doc_id}")
-
-
-def test_individual_exception_docs(docs):
-    """Test that individual exception documents exist."""
-    print("\nTesting individual exception documents...")
-
-    # These are the individual exception docs we configured
-    exception_ids = [
-        'sales-proposal-001',
-        'hr-policy-002',
-        'engineering-guide-006',
-    ]
-
-    for doc_id in exception_ids:
-        found = any(d['doc_id'] == doc_id for d in docs)
-        assert found, f"Individual exception document not found: {doc_id}"
-        print(f"  ✅ Found exception doc: {doc_id}")
-
-
-def test_doc_id_format(docs):
-    """Test that all doc_ids follow the naming convention."""
-    print("\nTesting doc_id format...")
-
-    import re
-    pattern = re.compile(r'^[a-z]+-[a-z]+-\d{3}$')
-
-    for doc in docs:
-        doc_id = doc['doc_id']
-        assert pattern.match(doc_id), \
-            f"Invalid doc_id format: {doc_id} (expected: dept-category-###)"
-
-    print(f"  ✅ All {len(docs)} doc_ids follow naming convention")
-
-
-def main():
-    """Run all tests."""
-    print("=" * 60)
-    print("Dataset Integration Tests")
-    print("=" * 60)
-
-    try:
-        docs = test_document_loading()
-        test_department_distribution(docs)
-        test_document_structure(docs)
-        test_cross_department_docs(docs)
-        test_individual_exception_docs(docs)
-        test_doc_id_format(docs)
-
-        print("\n" + "=" * 60)
-        print("✅ All integration tests passed!")
-        print("=" * 60)
-        return True
-
-    except AssertionError as e:
-        print(f"\n❌ Test failed: {e}")
-        return False
-    except Exception as e:
-        print(f"\n❌ Error: {e}")
-        import traceback
-        traceback.print_exc()
-        return False
-
-
-if __name__ == "__main__":
-    success = main()
-    sys.exit(0 if success else 1)
diff --git a/agentic-rag-authorization/test_improvements.py b/agentic-rag-authorization/test_improvements.py
deleted file mode 100644
index 831398b..0000000
--- a/agentic-rag-authorization/test_improvements.py
+++ /dev/null
@@ -1,277 +0,0 @@
-"""
-Test script to verify production improvements.
-
-This tests the new features without requiring running services.
-"""
-
-import json
-import sys
-from io import StringIO
-
-
-def test_structured_logging():
-    """Test that logging produces valid JSON."""
-    print("\n=== Testing Structured Logging ===")
-
-    # Capture logging output
-    from agentic_rag.logging_config import setup_logging, get_logger
-
-    setup_logging(level="INFO")
-    logger = get_logger("test")
-
-    # Capture stdout
-    old_stdout = sys.stdout
-    sys.stdout = captured_output = StringIO()
-
-    # Log with context
-    logger.info(
-        "Test message",
-        extra={
-            "subject_id": "alice",
-            "document_count": 5,
-            "duration_ms": 123.45,
-        },
-    )
-
-    # Restore stdout
-    sys.stdout = old_stdout
-    output = captured_output.getvalue()
-
-    # Parse as JSON
-    try:
-        log_entry = json.loads(output.strip())
-        assert log_entry["level"] == "INFO"
-        assert log_entry["logger"] == "agentic_rag.test"
-        assert log_entry["message"] == "Test message"
-        assert log_entry["subject_id"] == "alice"
-        assert log_entry["document_count"] == 5
-        assert log_entry["duration_ms"] == 123.45
-        print("✅ Structured logging works correctly")
-        print(f"   Sample log: {json.dumps(log_entry, indent=2)[:200]}...")
-        return True
-    except (json.JSONDecodeError, AssertionError, KeyError) as e:
-        print(f"❌ Structured logging failed: {e}")
-        print(f"   Output: {output}")
-        return False
-
-
-def test_connection_pooling():
-    """Test that singletons are created correctly."""
-    print("\n=== Testing Connection Pooling ===")
-
-    try:
-        from agentic_rag.grpc_helpers import (
-            get_spicedb_client,
-            reset_spicedb_client,
-            _spicedb_client,
-        )
-        from agentic_rag.elasticsearch_client import (
-            get_elasticsearch_client,
-            reset_elasticsearch_client,
-            _es_client,
-        )
-
-        # Verify reset functions exist
-        assert callable(reset_spicedb_client)
-        assert callable(reset_elasticsearch_client)
-
-        print("✅ Connection pooling functions defined correctly")
-        print("   - get_spicedb_client() available")
-        print("   - get_elasticsearch_client() available")
-        print("   - reset_*_client() functions available")
-        return True
-    except (ImportError, AssertionError) as e:
-        print(f"❌ Connection pooling failed: {e}")
-        return False
-
-
-def test_batch_permissions():
-    """Test that langchain-spicedb SpiceDBAuthorizer is importable."""
-    print("\n=== Testing Batch Permission Checker (langchain-spicedb) ===")
-
-    try:
-        from langchain_spicedb.core import SpiceDBAuthorizer
-        import inspect
-
-        sig = inspect.signature(SpiceDBAuthorizer.filter_documents)
-        params = list(sig.parameters.keys())
-
-        assert "subject_id" in params
-        assert "documents" in params
-
-        print("✅ SpiceDBAuthorizer.filter_documents available from langchain-spicedb")
-        print(f"   Method signature: filter_documents{sig}")
-        return True
-    except (ImportError, AssertionError) as e:
-        print(f"❌ SpiceDBAuthorizer check failed: {e}")
-        return False
-
-
-def test_validation():
-    """Test input validation."""
-    print("\n=== Testing Input Validation ===")
-
-    try:
-        from agentic_rag.validation import (
-            validate_query,
-            validate_subject_id,
-            ValidationError,
-        )
-
-        # Test valid inputs
-        query = validate_query("  test query  ")
-        assert query == "test query"
-
-        subject = validate_subject_id("alice_123")
-        assert subject == "alice_123"
-
-        # Test invalid inputs
-        try:
-            validate_query("")
-            print("❌ Empty query should raise ValidationError")
-            return False
-        except ValidationError:
-            pass  # Expected
-
-        try:
-            validate_subject_id("alice@example.com")
-            print("❌ Invalid characters should raise ValidationError")
-            return False
-        except ValidationError:
-            pass  # Expected
-
-        # Test truncation
-        long_query = "x" * 2000
-        truncated = validate_query(long_query)
-        assert len(truncated) == 1000
-
-        print("✅ Input validation works correctly")
-        print("   - Empty query rejected")
-        print("   - Invalid characters rejected")
-        print("   - Long queries truncated")
-        return True
-    except Exception as e:
-        print(f"❌ Input validation failed: {e}")
-        return False
-
-
-def test_error_handling():
-    """Test that error handling is present in nodes."""
-    print("\n=== Testing Error Handling ===")
-
-    try:
-        # Check retrieval node has try-except
-        with open("agentic_rag/nodes/retrieval_node.py", "r") as f:
-            retrieval_code = f.read()
-
-        assert "except Exception as e:" in retrieval_code
-        assert "logger.error" in retrieval_code
-
-        # Check authorization node has logging
-        with open("agentic_rag/nodes/authorization_node.py", "r") as f:
-            auth_code = f.read()
-
-        assert "logger.info" in auth_code
-        assert "SpiceDBAuthorizer" in auth_code
-
-        print("✅ Error handling implemented correctly")
-        print("   - Retrieval node has try-except")
-        print("   - Authorization node uses langchain-spicedb SpiceDBAuthorizer")
-        print("   - Errors logged with logger.error")
-        return True
-    except Exception as e:
-        print(f"❌ Error handling check failed: {e}")
-        return False
-
-
-def test_config():
-    """Test that config has log_level field."""
-    print("\n=== Testing Configuration ===")
-
-    try:
-        from agentic_rag.config import get_config
-
-        config = get_config()
-
-        assert hasattr(config, "log_level")
-        assert config.log_level == "INFO"  # Default value
-
-        print("✅ Configuration updated correctly")
-        print(f"   - log_level field added (default: {config.log_level})")
-        return True
-    except Exception as e:
-        print(f"❌ Configuration check failed: {e}")
-        return False
-
-
-def test_graph_validation():
-    """Test that graph has validation wrapper."""
-    print("\n=== Testing Graph Validation ===")
-
-    try:
-        from agentic_rag.graph import run_agentic_rag
-        import inspect
-
-        # Verify function signature
-        sig = inspect.signature(run_agentic_rag)
-        params = list(sig.parameters.keys())
-
-        assert "query" in params
-        assert "subject_id" in params
-
-        print("✅ Graph validation wrapper defined correctly")
-        print(f"   Function: run_agentic_rag{sig}")
-        return True
-    except Exception as e:
-        print(f"❌ Graph validation check failed: {e}")
-        return False
-
-
-def main():
-    """Run all tests."""
-    print("\n" + "=" * 60)
-    print("PRODUCTION IMPROVEMENTS VERIFICATION")
-    print("=" * 60)
-
-    tests = [
-        test_structured_logging,
-        test_connection_pooling,
-        test_batch_permissions,
-        test_validation,
-        test_error_handling,
-        test_config,
-        test_graph_validation,
-    ]
-
-    results = []
-    for test in tests:
-        try:
-            result = test()
-            results.append(result)
-        except Exception as e:
-            print(f"\n❌ Test {test.__name__} crashed: {e}")
-            import traceback
-
-            traceback.print_exc()
-            results.append(False)
-
-    # Summary
-    print("\n" + "=" * 60)
-    print("SUMMARY")
-    print("=" * 60)
-
-    passed = sum(results)
-    total = len(results)
-
-    print(f"\nPassed: {passed}/{total}")
-
-    if passed == total:
-        print("\n✅ All production improvements verified successfully!")
-        return 0
-    else:
-        print(f"\n⚠️  {total - passed} tests failed")
-        return 1
-
-
-if __name__ == "__main__":
-    sys.exit(main())
diff --git a/agentic-rag-authorization/tests/conftest.py b/agentic-rag-authorization/tests/conftest.py
index 6819429..65f7182 100644
--- a/agentic-rag-authorization/tests/conftest.py
+++ b/agentic-rag-authorization/tests/conftest.py
@@ -1,25 +1,82 @@
-"""Test fixtures for agentic RAG tests."""
+"""Shared fixtures and service-reachability helpers for the test suite.
+
+Deterministic tests mock Elasticsearch/SpiceDB/Mistral and run with no infra.
+Tests marked ``@pytest.mark.integration`` exercise real services; they use the
+``require_services`` fixture to skip automatically when a service is unreachable.
+Run the live suite with ``pytest -m integration`` after ``setup_environment.py``.
+"""
+
+import socket
+from urllib.parse import urlparse
 
 import pytest
 from elasticsearch import Elasticsearch
+
 from agentic_rag.grpc_helpers import create_insecure_spicedb_client
 
 
+def _tcp_reachable(host: str, port: int, timeout: float = 0.5) -> bool:
+    """Return True if a TCP connection to host:port succeeds quickly."""
+    try:
+        with socket.create_connection((host, port), timeout=timeout):
+            return True
+    except OSError:
+        return False
+
+
+def elasticsearch_reachable(url: str = "http://localhost:9200") -> bool:
+    parsed = urlparse(url)
+    return _tcp_reachable(parsed.hostname or "localhost", parsed.port or 9200)
+
+
+def spicedb_reachable(endpoint: str = "localhost:50051") -> bool:
+    host, _, port = endpoint.partition(":")
+    return _tcp_reachable(host or "localhost", int(port or "50051"))
+
+
+def mistral_configured() -> bool:
+    from agentic_rag.config import get_config
+
+    return bool(get_config().mistral_api_key)
+
+
+@pytest.fixture
+def require_services():
+    """Skip the calling test unless the requested live services are available.
+
+    Reachability is checked against the SAME endpoints the code connects to
+    (``config.elasticsearch_url`` / ``config.spicedb_endpoint``), so pointing
+    those env vars at a non-localhost host gates the tests correctly.
+    """
+    from agentic_rag.config import get_config
+
+    def _require(*, elasticsearch: bool = False, spicedb: bool = False, mistral: bool = False):
+        config = get_config()
+        if elasticsearch and not elasticsearch_reachable(config.elasticsearch_url):
+            pytest.skip(f"Elasticsearch not reachable at {config.elasticsearch_url}")
+        if spicedb and not spicedb_reachable(config.spicedb_endpoint):
+            pytest.skip(f"SpiceDB not reachable at {config.spicedb_endpoint}")
+        if mistral and not mistral_configured():
+            pytest.skip("MISTRAL_API_KEY not configured")
+
+    return _require
+
+
 @pytest.fixture
 def elasticsearch_client():
-    """Create Elasticsearch client for tests."""
+    """Create Elasticsearch client for integration tests."""
     return Elasticsearch("http://localhost:9200")
 
 
 @pytest.fixture
 def spicedb_client():
-    """Create SpiceDB client for tests."""
+    """Create SpiceDB client for integration tests."""
     return create_insecure_spicedb_client("localhost:50051", "devtoken")
 
 
 @pytest.fixture
 def sample_state():
-    """Create sample state for testing."""
+    """Create a sample graph state for testing."""
     return {
         "query": "Test query",
         "subject_id": "alice",
diff --git a/agentic-rag-authorization/tests/test_authorization_integration.py b/agentic-rag-authorization/tests/test_authorization_integration.py
new file mode 100644
index 0000000..0655b38
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_authorization_integration.py
@@ -0,0 +1,55 @@
+"""Live integration test for authorization against a real SpiceDB.
+
+Requires SpiceDB running with the schema and relationships created by
+``examples/setup_environment.py``. Skipped automatically when SpiceDB is
+unreachable. Run with ``pytest -m integration``.
+"""
+
+import pytest
+from langchain_core.documents import Document
+from langchain_spicedb.core import SpiceDBAuthorizer
+
+from agentic_rag.config import get_config
+
+
+def _doc(doc_id, dept):
+    return Document(
+        page_content=f"content of {doc_id}",
+        metadata={"doc_id": doc_id, "department": dept},
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_mixed_authorization_against_live_spicedb(require_services):
+    """alice (engineering) can view engineering docs but not HR docs: a mixed result.
+
+    This is the "some but not all" case against real SpiceDB. We use fixed,
+    known document IDs rather than semantic retrieval so the authorized/denied
+    split is deterministic.
+    """
+    require_services(spicedb=True)
+
+    config = get_config()
+    authorizer = SpiceDBAuthorizer(
+        spicedb_endpoint=config.spicedb_endpoint,
+        spicedb_token=config.spicedb_token,
+        resource_type="document",
+        subject_type="user",
+        permission="view",
+        resource_id_key="doc_id",
+    )
+
+    documents = [
+        _doc("engineering-architecture-001", "engineering"),  # alice CAN view
+        _doc("hr-policy-003", "hr"),                           # alice CANNOT view
+    ]
+
+    result = await authorizer.filter_documents(documents=documents, subject_id="alice")
+
+    authorized_ids = {d.metadata["doc_id"] for d in result.authorized_documents}
+    assert "engineering-architecture-001" in authorized_ids
+    assert "hr-policy-003" not in authorized_ids
+    assert "hr-policy-003" in result.denied_resource_ids
+    assert result.total_authorized == 1
+    assert result.total_retrieved == 2
diff --git a/agentic-rag-authorization/tests/test_authorization_node.py b/agentic-rag-authorization/tests/test_authorization_node.py
new file mode 100644
index 0000000..a49cf7a
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_authorization_node.py
@@ -0,0 +1,134 @@
+"""Unit tests for the authorization node (the security boundary), SpiceDB mocked.
+
+The authorization node ALWAYS runs and filters retrieved documents through
+SpiceDB. These tests mock ``SpiceDBAuthorizer`` so the node's state-handling
+logic is covered deterministically, with no live SpiceDB. For a real round-trip
+against SpiceDB, see ``test_authorization_integration.py``.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from langchain_core.documents import Document
+from langchain_spicedb.core import AuthorizationResult
+
+from agentic_rag.nodes.authorization_node import authorization_node
+
+
+def _doc(doc_id, content="content", title="Title", dept="engineering"):
+    return Document(
+        page_content=content,
+        metadata={
+            "doc_id": doc_id,
+            "title": title,
+            "department": dept,
+            "classification": "internal",
+        },
+    )
+
+
+def _state(documents, subject_id="alice"):
+    return {
+        "query": "q",
+        "subject_id": subject_id,
+        "max_attempts": 1,
+        "retrieval_attempt": 1,
+        "messages": [],
+        "reasoning": [],
+        "retrieved_documents": documents,
+        "authorized_documents": [],
+        "denied_count": 0,
+        "authorization_passed": False,
+        "answer": "",
+    }
+
+
+def _patch_authorizer(result):
+    """Patch the module-level authorizer to return ``result`` from filter_documents."""
+    authorizer = MagicMock()
+    authorizer.filter_documents = AsyncMock(return_value=result)
+    return patch(
+        "agentic_rag.nodes.authorization_node._get_authorizer",
+        return_value=authorizer,
+    ), authorizer
+
+
+@pytest.mark.asyncio
+async def test_all_documents_authorized():
+    docs = [_doc("doc-001"), _doc("doc-002")]
+    result = AuthorizationResult(
+        authorized_documents=docs,
+        total_retrieved=2,
+        total_authorized=2,
+        authorization_rate=1.0,
+        denied_resource_ids=[],
+    )
+    cm, authorizer = _patch_authorizer(result)
+    with cm:
+        out = await authorization_node(_state(docs))
+
+    assert out["authorization_passed"] is True
+    assert out["denied_count"] == 0
+    assert out["authorized_documents"] == docs
+    authorizer.filter_documents.assert_awaited_once_with(documents=docs, subject_id="alice")
+
+
+@pytest.mark.asyncio
+async def test_no_documents_authorized():
+    docs = [_doc("doc-001"), _doc("doc-002")]
+    result = AuthorizationResult(
+        authorized_documents=[],
+        total_retrieved=2,
+        total_authorized=0,
+        authorization_rate=0.0,
+        denied_resource_ids=["doc-001", "doc-002"],
+    )
+    cm, _ = _patch_authorizer(result)
+    with cm:
+        out = await authorization_node(_state(docs))
+
+    assert out["authorization_passed"] is False
+    assert out["denied_count"] == 2
+    assert out["authorized_documents"] == []
+
+
+@pytest.mark.asyncio
+async def test_mixed_authorization_some_but_not_all():
+    """User can view some retrieved documents but not others (the common case)."""
+    authorized = _doc("engineering-guide-001", dept="engineering")
+    denied = _doc("hr-policy-003", dept="hr")
+    result = AuthorizationResult(
+        authorized_documents=[authorized],
+        total_retrieved=2,
+        total_authorized=1,
+        authorization_rate=0.5,
+        denied_resource_ids=["hr-policy-003"],
+    )
+    cm, _ = _patch_authorizer(result)
+    with cm:
+        out = await authorization_node(_state([authorized, denied]))
+
+    # At least one authorized document means the flow may proceed to generation,
+    # but the partial denial must still be recorded.
+    assert out["authorization_passed"] is True
+    assert out["denied_count"] == 1
+    assert out["authorized_documents"] == [authorized]
+    # The denied document must never leak into the authorized set.
+    assert denied not in out["authorized_documents"]
+
+
+@pytest.mark.asyncio
+async def test_empty_retrieved_documents():
+    result = AuthorizationResult(
+        authorized_documents=[],
+        total_retrieved=0,
+        total_authorized=0,
+        authorization_rate=0.0,
+        denied_resource_ids=[],
+    )
+    cm, authorizer = _patch_authorizer(result)
+    with cm:
+        out = await authorization_node(_state([]))
+
+    assert out["authorization_passed"] is False
+    assert out["denied_count"] == 0
+    authorizer.filter_documents.assert_awaited_once_with(documents=[], subject_id="alice")
diff --git a/agentic-rag-authorization/tests/test_basic_flow.py b/agentic-rag-authorization/tests/test_basic_flow.py
index 8a7aea8..1bf14a5 100644
--- a/agentic-rag-authorization/tests/test_basic_flow.py
+++ b/agentic-rag-authorization/tests/test_basic_flow.py
@@ -1,86 +1,75 @@
-"""Basic integration tests for agentic RAG."""
+"""End-to-end integration tests for the agentic RAG graph.
+
+These run the full graph against live Elasticsearch + SpiceDB + Mistral, so they
+are marked ``integration`` and skip automatically when those services are
+unavailable. Run with ``pytest -m integration`` after seeding the environment
+via ``examples/setup_environment.py``.
+"""
 
 import pytest
 from agentic_rag.graph import build_agentic_rag_graph
 
 
+def _initial_state(query, subject_id, max_attempts):
+    return {
+        "query": query,
+        "subject_id": subject_id,
+        "max_attempts": max_attempts,
+        "retrieval_attempt": 0,
+        "messages": [],
+        "reasoning": [],
+        "retrieved_documents": [],
+        "authorized_documents": [],
+        "denied_count": 0,
+        "authorization_passed": False,
+        "answer": "",
+    }
+
+
+@pytest.mark.integration
 @pytest.mark.asyncio
-async def test_authorized_access():
-    """Test that authorized user can access documents."""
+async def test_authorized_access(require_services):
+    """Authorized user (alice, engineering) can access engineering documents."""
+    require_services(elasticsearch=True, spicedb=True, mistral=True)
     graph = build_agentic_rag_graph()
 
     result = await graph.ainvoke(
-        {
-            "query": "What are our engineering best practices?",
-            "subject_id": "alice",
-            "max_attempts": 3,
-            "retrieval_attempt": 0,
-            "messages": [],
-            "reasoning": [],
-            "retrieved_documents": [],
-            "authorized_documents": [],
-            "denied_count": 0,
-            "authorization_passed": False,
-            "answer": "",
-        }
+        _initial_state("What are our engineering best practices?", "alice", 3)
     )
 
-    # Alice should have access to engineering documents
     assert len(result["authorized_documents"]) > 0
     assert result["authorization_passed"] is True
     assert result["answer"] != ""
 
 
+@pytest.mark.integration
 @pytest.mark.asyncio
-async def test_denied_access():
-    """Test that unauthorized user is denied access."""
+async def test_denied_access(require_services):
+    """Unauthorized user (bob, sales) is denied access to HR documents."""
+    require_services(elasticsearch=True, spicedb=True, mistral=True)
     graph = build_agentic_rag_graph()
 
     result = await graph.ainvoke(
-        {
-            "query": "What are the HR policies and employee guidelines?",
-            "subject_id": "bob",
-            "max_attempts": 3,
-            "retrieval_attempt": 0,
-            "messages": [],
-            "reasoning": [],
-            "retrieved_documents": [],
-            "authorized_documents": [],
-            "denied_count": 0,
-            "authorization_passed": False,
-            "answer": "",
-        }
+        _initial_state("What are the HR policies and employee guidelines?", "bob", 3)
     )
 
-    # Bob (sales) should not have access to HR documents
     assert len(result["authorized_documents"]) == 0
     assert result["denied_count"] > 0
     assert result["answer"] != ""
-    # Answer should explain the access denial
+    # Answer should explain the access denial.
     assert "access" in result["answer"].lower() or "permission" in result["answer"].lower()
 
 
+@pytest.mark.integration
 @pytest.mark.asyncio
-async def test_retrieval_attempts():
-    """Test that system respects max retrieval attempts."""
+async def test_retrieval_attempts(require_services):
+    """System respects max retrieval attempts when authorization keeps failing."""
+    require_services(elasticsearch=True, spicedb=True, mistral=True)
     graph = build_agentic_rag_graph()
 
     result = await graph.ainvoke(
-        {
-            "query": "Query that will fail authorization",
-            "subject_id": "unauthorized_user",
-            "max_attempts": 2,
-            "retrieval_attempt": 0,
-            "messages": [],
-            "reasoning": [],
-            "retrieved_documents": [],
-            "authorized_documents": [],
-            "denied_count": 0,
-            "authorization_passed": False,
-            "answer": "",
-        }
+        _initial_state("Query that will fail authorization", "unauthorized_user", 2)
     )
 
-    # Should not exceed max attempts
     assert result["retrieval_attempt"] <= 2
     assert result["answer"] != ""
diff --git a/agentic-rag-authorization/tests/test_config.py b/agentic-rag-authorization/tests/test_config.py
index 90beb81..2fca9bb 100644
--- a/agentic-rag-authorization/tests/test_config.py
+++ b/agentic-rag-authorization/tests/test_config.py
@@ -38,12 +38,3 @@ def test_config_elasticsearch_defaults():
         config = Config.from_env()
     assert config.elasticsearch_url == "http://localhost:9200"
     assert config.elasticsearch_api_key == ""
-
-
-def test_config_has_no_milvus_weaviate_or_openai_fields():
-    config = Config.from_env()
-    assert not hasattr(config, "milvus_uri")
-    assert not hasattr(config, "milvus_token")
-    assert not hasattr(config, "weaviate_url")
-    assert not hasattr(config, "weaviate_api_key")
-    assert not hasattr(config, "openai_api_key")
diff --git a/agentic-rag-authorization/tests/test_dataset.py b/agentic-rag-authorization/tests/test_dataset.py
new file mode 100644
index 0000000..07035ef
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_dataset.py
@@ -0,0 +1,70 @@
+"""Tests for the realistic document dataset in data/documents.
+
+(Moved from the repo-root ``test_dataset_integration.py`` and converted to
+idiomatic pytest: a shared ``docs`` fixture replaces the hand-threaded
+function arguments and ``main()`` runner.)
+"""
+
+import re
+import sys
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT / "scripts"))
+
+from parse_documents import load_all_documents  # noqa: E402
+
+
+@pytest.fixture(scope="module")
+def docs():
+    return load_all_documents(str(REPO_ROOT / "data" / "documents"))
+
+
+def test_loads_all_documents_with_required_fields(docs):
+    assert len(docs) == 50
+    for doc in docs:
+        for field in ("doc_id", "title", "content", "department", "classification"):
+            assert field in doc, f"Missing {field} in {doc.get('doc_id', doc)}"
+
+
+def test_department_distribution(docs):
+    counts = {}
+    for doc in docs:
+        counts[doc["department"]] = counts.get(doc["department"], 0) + 1
+    assert counts == {
+        "engineering": 15,
+        "sales": 10,
+        "hr": 10,
+        "finance": 10,
+        "public": 5,
+    }
+
+
+def test_public_documents_classified_public(docs):
+    public = [d for d in docs if d["department"] == "public"]
+    assert len(public) == 5
+    assert all(d["classification"] == "public" for d in public)
+
+
+@pytest.mark.parametrize(
+    "doc_id",
+    ["engineering-architecture-001", "sales-guide-005", "hr-policy-001"],
+)
+def test_cross_department_documents_exist(docs, doc_id):
+    assert any(d["doc_id"] == doc_id for d in docs)
+
+
+@pytest.mark.parametrize(
+    "doc_id",
+    ["sales-proposal-001", "hr-policy-002", "engineering-guide-006"],
+)
+def test_individual_exception_documents_exist(docs, doc_id):
+    assert any(d["doc_id"] == doc_id for d in docs)
+
+
+def test_doc_id_naming_convention(docs):
+    pattern = re.compile(r"^[a-z]+-[a-z]+-\d{3}$")
+    for doc in docs:
+        assert pattern.match(doc["doc_id"]), f"Invalid doc_id: {doc['doc_id']}"
diff --git a/agentic-rag-authorization/tests/test_elasticsearch_integration.py b/agentic-rag-authorization/tests/test_elasticsearch_integration.py
new file mode 100644
index 0000000..04c3558
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_elasticsearch_integration.py
@@ -0,0 +1,54 @@
+"""Live integration test that pulls documents from a real Elasticsearch.
+
+Requires Elasticsearch running and indexed by ``examples/setup_environment.py``,
+plus a configured ``MISTRAL_API_KEY`` for the query embedding. Skipped
+automatically when those are unavailable. Run with ``pytest -m integration``.
+"""
+
+import pytest
+
+from agentic_rag.elasticsearch_client import reset_elasticsearch_client
+from agentic_rag.nodes.retrieval_node import retrieval_node
+
+
+@pytest.fixture(autouse=True)
+def _reset_es_singleton():
+    """Isolate the module-level Elasticsearch client around each test."""
+    reset_elasticsearch_client()
+    yield
+    reset_elasticsearch_client()
+
+
+@pytest.mark.integration
+def test_retrieval_pulls_documents_from_elasticsearch(require_services):
+    require_services(elasticsearch=True, mistral=True)
+
+    state = {
+        "query": "What are our engineering best practices?",
+        "subject_id": "alice",
+        "max_attempts": 1,
+        "retrieval_attempt": 0,
+        "messages": [],
+        "reasoning": [],
+        "retrieved_documents": [],
+        "authorized_documents": [],
+        "denied_count": 0,
+        "authorization_passed": False,
+        "answer": "",
+    }
+
+    out = retrieval_node(state)
+
+    # retrieval_node swallows exceptions and returns []; surface any such error
+    # (bad embedding key, missing index, ...) instead of a vague count assertion.
+    last_message = out["messages"][-1].content if out["messages"] else ""
+    assert "failed" not in last_message.lower(), f"retrieval_node reported an error: {last_message}"
+
+    docs = out["retrieved_documents"]
+    assert len(docs) > 0, "Elasticsearch returned no documents — run examples/setup_environment.py first"
+    assert out["retrieval_attempt"] == 1
+    for doc in docs:
+        assert doc.metadata.get("doc_id")
+        assert "department" in doc.metadata
+        assert "classification" in doc.metadata
+        assert doc.page_content
diff --git a/agentic-rag-authorization/tests/test_generation_node.py b/agentic-rag-authorization/tests/test_generation_node.py
new file mode 100644
index 0000000..e2b5e98
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_generation_node.py
@@ -0,0 +1,88 @@
+"""Unit tests for the generation node, focused on preventing document leaks.
+
+The security guarantee is architectural: the authorization node filters denied
+documents out of state BEFORE generation runs, so generation only ever sees
+``state["authorized_documents"]``. These tests assert that denied content never
+reaches the LLM prompt (and therefore cannot appear in the generated answer).
+"""
+
+from unittest.mock import patch
+
+from langchain_core.documents import Document
+from langchain_core.messages import AIMessage
+from langchain_core.runnables import RunnableLambda
+
+from agentic_rag.nodes.generation_node import generation_node
+
+AUTHORIZED_CONTENT = "ENGINEERING_BEST_PRACTICES_AUTHORIZED"
+AUTHORIZED_TITLE = "Engineering Guide"
+DENIED_CONTENT = "CONFIDENTIAL_HR_SALARY_DATA_DENIED"
+
+
+def _state(authorized_documents, denied_count, retrieved_documents=None):
+    return {
+        "query": "What are our engineering best practices?",
+        "subject_id": "alice",
+        "max_attempts": 1,
+        "retrieval_attempt": 1,
+        "messages": [],
+        "reasoning": [],
+        "retrieved_documents": retrieved_documents or [],
+        "authorized_documents": authorized_documents,
+        "denied_count": denied_count,
+        "authorization_passed": bool(authorized_documents),
+        "answer": "",
+    }
+
+
+def _capturing_llm(captured):
+    """A fake LLM (Runnable) that records the prompt it is handed."""
+
+    def _invoke(prompt_value):
+        captured["prompt"] = prompt_value.to_string()
+        return AIMessage(content="Answer based only on authorized documents.")
+
+    return RunnableLambda(_invoke)
+
+
+def test_denied_document_content_never_reaches_prompt_or_answer():
+    captured = {}
+    authorized = Document(
+        page_content=AUTHORIZED_CONTENT,
+        metadata={"title": AUTHORIZED_TITLE, "doc_id": "engineering-guide-001"},
+    )
+    # The denied document was retrieved but filtered out by authorization: it is
+    # present in retrieved_documents but NOT in authorized_documents. If generation
+    # ever read from retrieved_documents, this content would leak into the prompt.
+    denied = Document(
+        page_content=DENIED_CONTENT,
+        metadata={"title": "Confidential HR Salaries", "doc_id": "hr-policy-003"},
+    )
+    state = _state([authorized], denied_count=1, retrieved_documents=[authorized, denied])
+
+    with patch(
+        "agentic_rag.nodes.generation_node.get_llm",
+        return_value=_capturing_llm(captured),
+    ):
+        out = generation_node(state)
+
+    # Authorized content IS present in the prompt...
+    assert AUTHORIZED_CONTENT in captured["prompt"]
+    assert AUTHORIZED_TITLE in captured["prompt"]
+    # ...and the denied document's content — though present in retrieved_documents —
+    # is absent from both the prompt and the answer.
+    assert DENIED_CONTENT not in captured["prompt"]
+    assert DENIED_CONTENT not in out["answer"]
+
+
+def test_no_authorized_documents_produces_no_document_content():
+    captured = {}
+    with patch(
+        "agentic_rag.nodes.generation_node.get_llm",
+        return_value=_capturing_llm(captured),
+    ):
+        out = generation_node(_state([], denied_count=2))
+
+    assert "No authorized documents available." in captured["prompt"]
+    assert DENIED_CONTENT not in captured["prompt"]
+    assert out["answer"]
diff --git a/agentic-rag-authorization/tests/test_graph_routing.py b/agentic-rag-authorization/tests/test_graph_routing.py
new file mode 100644
index 0000000..c7cc3dc
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_graph_routing.py
@@ -0,0 +1,56 @@
+"""Unit tests for the graph routing / conditional-edge functions (graph.py:14-51).
+
+These pure functions decide the control flow of the state machine after
+authorization and after reasoning. They are deterministic, so no mocking or
+live services are required.
+"""
+
+from agentic_rag.graph import should_retry_or_generate, should_reason_or_generate
+
+
+def _state(**overrides):
+    state = {
+        "retrieval_attempt": 0,
+        "max_attempts": 1,
+        "authorized_documents": [],
+        "authorization_passed": False,
+    }
+    state.update(overrides)
+    return state
+
+
+class TestShouldRetryOrGenerate:
+    def test_retries_when_attempts_remain_and_no_authorized_docs(self):
+        state = _state(retrieval_attempt=1, max_attempts=3, authorized_documents=[])
+        assert should_retry_or_generate(state) == "retrieve"
+
+    def test_generates_when_authorized_docs_exist(self):
+        state = _state(retrieval_attempt=1, max_attempts=3, authorized_documents=["doc"])
+        assert should_retry_or_generate(state) == "generate"
+
+    def test_generates_when_attempts_exhausted(self):
+        state = _state(retrieval_attempt=3, max_attempts=3, authorized_documents=[])
+        assert should_retry_or_generate(state) == "generate"
+
+    def test_generates_when_exhausted_even_with_docs(self):
+        state = _state(retrieval_attempt=3, max_attempts=3, authorized_documents=["doc"])
+        assert should_retry_or_generate(state) == "generate"
+
+
+class TestShouldReasonOrGenerate:
+    def test_generates_when_authorization_passed(self):
+        state = _state(authorization_passed=True, max_attempts=3, retrieval_attempt=1)
+        assert should_reason_or_generate(state) == "generate"
+
+    def test_reasons_when_denied_and_attempts_remain(self):
+        state = _state(authorization_passed=False, max_attempts=3, retrieval_attempt=1)
+        assert should_reason_or_generate(state) == "reason"
+
+    def test_generates_when_denied_and_single_attempt(self):
+        # Default max_attempts=1 skips reasoning entirely for efficiency.
+        state = _state(authorization_passed=False, max_attempts=1, retrieval_attempt=1)
+        assert should_reason_or_generate(state) == "generate"
+
+    def test_generates_when_denied_and_attempts_exhausted(self):
+        state = _state(authorization_passed=False, max_attempts=2, retrieval_attempt=2)
+        assert should_reason_or_generate(state) == "generate"
diff --git a/agentic-rag-authorization/tests/test_improvements.py b/agentic-rag-authorization/tests/test_improvements.py
new file mode 100644
index 0000000..5017c42
--- /dev/null
+++ b/agentic-rag-authorization/tests/test_improvements.py
@@ -0,0 +1,98 @@
+"""Tests for production-hardening features: structured logging, connection
+pooling, input validation, and configuration.
+
+(Moved from the repo-root ``test_improvements.py`` and converted from a
+print/``main()`` script into real pytest assertions.)
+"""
+
+import json
+import logging
+
+import pytest
+
+
+def test_structured_logging_emits_json(capsys):
+    from agentic_rag.logging_config import setup_logging, get_logger
+
+    root = logging.getLogger()
+    saved_handlers, saved_level = root.handlers[:], root.level
+    try:
+        setup_logging(level="INFO")
+        logger = get_logger("improvements_test")
+        logger.info(
+            "Test message",
+            extra={"subject_id": "alice", "document_count": 5, "duration_ms": 123.45},
+        )
+        lines = [ln for ln in capsys.readouterr().out.splitlines() if ln.strip()]
+        entry = json.loads(lines[-1])
+    finally:
+        root.handlers, root.level = saved_handlers, saved_level
+
+    assert entry["level"] == "INFO"
+    assert entry["logger"] == "agentic_rag.improvements_test"
+    assert entry["message"] == "Test message"
+    assert entry["subject_id"] == "alice"
+    assert entry["document_count"] == 5
+    assert entry["duration_ms"] == 123.45
+
+
+def test_connection_pooling_helpers_exist():
+    from agentic_rag.grpc_helpers import get_spicedb_client, reset_spicedb_client
+    from agentic_rag.elasticsearch_client import (
+        get_elasticsearch_client,
+        reset_elasticsearch_client,
+    )
+
+    assert callable(get_spicedb_client)
+    assert callable(reset_spicedb_client)
+    assert callable(get_elasticsearch_client)
+    assert callable(reset_elasticsearch_client)
+
+
+def test_spicedb_authorizer_filter_documents_signature():
+    import inspect
+    from langchain_spicedb.core import SpiceDBAuthorizer
+
+    params = inspect.signature(SpiceDBAuthorizer.filter_documents).parameters
+    assert "subject_id" in params
+    assert "documents" in params
+
+
+def test_input_validation():
+    from agentic_rag.validation import (
+        validate_query,
+        validate_subject_id,
+        ValidationError,
+    )
+
+    assert validate_query("  test query  ") == "test query"
+    assert validate_subject_id("alice_123") == "alice_123"
+
+    with pytest.raises(ValidationError):
+        validate_query("")
+    with pytest.raises(ValidationError):
+        validate_subject_id("alice@example.com")
+
+    assert len(validate_query("x" * 2000)) == 1000
+
+
+# Node error-handling is covered behaviorally rather than by source-text grep:
+#   - retrieval  -> test_retrieval_node.py::test_retrieval_node_increments_attempt_on_failure
+#   - authorization -> test_authorization_node.py
+
+
+def test_config_has_log_level():
+    from agentic_rag.config import get_config
+
+    config = get_config()
+    assert hasattr(config, "log_level")
+    assert config.log_level == "INFO"
+
+
+def test_graph_entrypoint_signature():
+    import inspect
+    from agentic_rag.graph import run_agentic_rag
+
+    params = inspect.signature(run_agentic_rag).parameters
+    assert "query" in params
+    assert "subject_id" in params
diff --git a/agentic-rag-authorization/tests/test_retrieval_node.py b/agentic-rag-authorization/tests/test_retrieval_node.py
index 38fe6c5..81f1abe 100644
--- a/agentic-rag-authorization/tests/test_retrieval_node.py
+++ b/agentic-rag-authorization/tests/test_retrieval_node.py
@@ -56,6 +56,23 @@ def test_retrieval_node_returns_documents(sample_state):
     assert doc.page_content == "Content"
     assert result["retrieval_attempt"] == 1
 
+    # Exactly one Elasticsearch query is performed...
+    mock_client.search.assert_called_once()
+    # ...and it is the expected kNN query over the embedded user query.
+    _, search_kwargs = mock_client.search.call_args
+    assert search_kwargs["index"] == "documents"
+    assert search_kwargs["knn"]["field"] == "embedding"
+    assert search_kwargs["knn"]["k"] == 5
+    assert search_kwargs["knn"]["num_candidates"] == 50
+    assert search_kwargs["knn"]["query_vector"] == [0.1] * 1024
+    assert search_kwargs["source"] == [
+        "doc_id",
+        "title",
+        "content",
+        "department",
+        "classification",
+    ]
+
 
 def test_retrieval_node_increments_attempt_on_failure(sample_state):
     from agentic_rag.nodes.retrieval_node import retrieval_node