diff --git a/.coveragerc b/.coveragerc
new file mode 100644
index 0000000..a693d91
--- /dev/null
+++ b/.coveragerc
@@ -0,0 +1,26 @@
+[run]
+source = chatbot_api/src, chatbot_frontend/src, hospital_neo4j_etl/src
+omit =
+    */tests/*
+    */test_*
+    */__init__.py
+    */conftest.py
+    */migrations/*
+    */venv/*
+    */.venv/*
+
+[report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    if self.debug:
+    if settings.DEBUG
+    raise AssertionError
+    raise NotImplementedError
+    if 0:
+    if __name__ == .__main__.:
+    class .*\bProtocol\):
+    @(abc\.)?abstractmethod
+
+[html]
+directory = htmlcov
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..b2419f9
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,142 @@
+name: CI
+
+on:
+  push:
+    branches: [master, main]
+  pull_request:
+    branches: [master, main]
+
+jobs:
+  code-quality:
+    name: Code Quality
+    runs-on: ubuntu-latest
+    if: ${{ !contains(github.event.head_commit.message, '[skip ci]') }}
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v3
+      with:
+        enable-cache: true
+        cache-dependency-glob: "**/pyproject.toml"
+
+    - name: Set up Python 3.11
+      run: uv python install 3.11
+
+    - name: Create virtual environment and install dependencies
+      run: |
+        uv venv --python 3.11
+        uv pip install -e "chatbot_api[dev]" -e "chatbot_frontend[dev]" -e "hospital_neo4j_etl[dev]"
+
+    - name: Run linting and formatting checks
+      run: |
+        uv run ruff check chatbot_api/src/ chatbot_frontend/src/ hospital_neo4j_etl/src/ tests/
+        uv run ruff format --check chatbot_api/src/ chatbot_frontend/src/ hospital_neo4j_etl/src/ tests/
+
+    - name: Run type checking
+      run: uv run mypy chatbot_api/src/
+
+  unit-tests:
+    name: Unit Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    if: ${{ !contains(github.event.head_commit.message, '[skip ci]') }}
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.11"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v3
+      with:
+        enable-cache: true
+        cache-dependency-glob: "**/pyproject.toml"
+
+    - name: Set up Python ${{ matrix.python-version }}
+      run: uv python install ${{ matrix.python-version }}
+
+    - name: Create virtual environment and install dependencies
+      run: |
+        uv venv --python ${{ matrix.python-version }}
+        uv pip install -e "chatbot_api[dev]" -e "chatbot_frontend[dev]" -e "hospital_neo4j_etl[dev]"
+
+    - name: Run unit tests with coverage
+      run: make test-unit
+
+  integration-tests:
+    name: Integration Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    if: ${{ !contains(github.event.head_commit.message, '[skip ci]') }}
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.11"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v3
+      with:
+        enable-cache: true
+        cache-dependency-glob: "**/pyproject.toml"
+
+    - name: Set up Python ${{ matrix.python-version }}
+      run: uv python install ${{ matrix.python-version }}
+
+    - name: Create virtual environment and install dependencies
+      run: |
+        uv venv --python ${{ matrix.python-version }}
+        uv pip install -e "chatbot_api[dev]" -e "chatbot_frontend[dev]" -e "hospital_neo4j_etl[dev]"
+
+    - name: Run integration tests with coverage
+      run: |
+        uv run pytest tests/integration/ -v --cov=chatbot_api/src --cov=chatbot_frontend/src --cov=hospital_neo4j_etl/src
+
+    - name: Upload coverage to Codecov
+      if: matrix.python-version == '3.11'  # Only upload once
+      uses: codecov/codecov-action@v3
+      with:
+        file: ./coverage.xml
+        fail_ci_if_error: false
+
+  build-validation:
+    name: Docker Build Validation
+    runs-on: ubuntu-latest
+    if: ${{ contains(github.event.head_commit.message, 'docker') || contains(github.event.head_commit.message, 'Docker') || github.event_name == 'pull_request' }}
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Check for Docker-related changes
+      uses: dorny/paths-filter@v2
+      id: docker-changes
+      with:
+        filters: |
+          docker:
+            - 'docker-compose.yml'
+            - '**/Dockerfile'
+            - '.dockerignore'
+
+    - name: Create dummy .env file for Docker build
+      if: steps.docker-changes.outputs.docker == 'true' || github.event_name == 'pull_request'
+      run: |
+        cat > .env << 'EOF'
+        # Dummy environment variables for CI build validation
+        OPENAI_API_KEY=sk-dummy-key-for-ci-build-only
+        NEO4J_URI=neo4j://localhost:7687
+        NEO4J_USERNAME=neo4j
+        NEO4J_PASSWORD=password
+        HOSPITAL_AGENT_MODEL=gpt-3.5-turbo
+        HOSPITAL_CYPHER_MODEL=gpt-3.5-turbo
+        HOSPITAL_QA_MODEL=gpt-3.5-turbo
+        CHATBOT_URL=http://localhost:8000/hospital-rag-agent
+        EOF
+
+    - name: Validate Docker Compose build
+      if: steps.docker-changes.outputs.docker == 'true' || github.event_name == 'pull_request'
+      run: |
+        docker compose build --no-cache
+        docker compose config --quiet
diff --git a/.gitignore b/.gitignore
index 6d6d89e..436e232 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,6 @@ __pycache__
 *.pyc
 .DS_Store
 
+*.egg-info
+*.coverage
+htmlcov/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..f3405af
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,23 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.1.8
+    hooks:
+      - id: ruff
+        args: [--fix, --exit-non-zero-on-fix]
+      - id: ruff-format
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.8.0
+    hooks:
+      - id: mypy
+        additional_dependencies: [types-requests]
+        args: [--ignore-missing-imports, --explicit-package-bases]
+        files: ^chatbot_api/
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
diff --git a/Makefile b/Makefile
index f80c400..76de0ed 100644
--- a/Makefile
+++ b/Makefile
@@ -8,8 +8,29 @@ start:
 stop:
 	docker compose down chatbot_api chatbot_frontend
 
-format:
+test:
+	uv run pytest tests/ -v --cov=chatbot_api/src --cov=chatbot_frontend/src --cov=hospital_neo4j_etl/src --ignore=tests/performance
+
+test-unit:
+	uv run pytest tests/unit/ -v
 
+test-integration:
+	uv run pytest tests/integration/ -v
+
+format:
+	uv run ruff format chatbot_api/src/ chatbot_frontend/src/ hospital_neo4j_etl/src/ tests/
 
 lint:
+	uv run ruff check chatbot_api/src/ chatbot_frontend/src/ hospital_neo4j_etl/src/ tests/
+
+format-check:
+	uv run ruff format --check chatbot_api/src/ chatbot_frontend/src/ hospital_neo4j_etl/src/ tests/
+
+lint-fix:
+	uv run ruff check --fix chatbot_api/src/ chatbot_frontend/src/ hospital_neo4j_etl/src/ tests/
+
+pre-commit:
+	uv run pre-commit run --all-files
 
+install-dev:
+	uv pip install -e "chatbot_api[dev]" -e "chatbot_frontend[dev]" -e "hospital_neo4j_etl[dev]"
diff --git a/README.md b/README.md
index a959789..93c3194 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,8 @@
-# Healthcare RAG Agent 
+# Healthcare RAG Agent
 
-[![Python](https://img.shields.io/badge/Python-3.8+-blue.svg)](https://www.python.org/)
+[![CI](https://github.com/asanmateu/medgraph-ai/workflows/CI/badge.svg)](https://github.com/asanmateu/medgraph-ai/actions)
+[![codecov](https://codecov.io/gh/asanmateu/medgraph-ai/branch/master/graph/badge.svg)](https://codecov.io/gh/asanmateu/medgraph-ai)
+[![Python](https://img.shields.io/badge/Python-3.9%20|%203.11-blue.svg)](https://www.python.org/)
 [![LangChain](https://img.shields.io/badge/LangChain-Latest-green.svg)](https://langchain.com/)
 [![Neo4j](https://img.shields.io/badge/Neo4j-5.0+-blue.svg)](https://neo4j.com/)
 [![Docker](https://img.shields.io/badge/Docker-Required-blue.svg)](https://www.docker.com/)
@@ -16,6 +18,7 @@ A Retrieval-Augmented Generation (RAG) agent designed for healthcare information
 - [Example Queries](#-example-queries)
 - [Database Design](#database-design)
 - [Technical Stack](#technical-stack)
+- [Contributing](#-contributing)
 - [Acknowledgments](#acknowledgments)
 
 ## 🎯 Overview
@@ -24,10 +27,10 @@ This project implements a healthcare-focused RAG chatbot that leverages LangChai
 
 ## ✨ Key Features
 
-✅ **Knowledge Graph Integration** - Neo4j for healthcare data relationships  
-✅ **RESTful API** - FastAPI-powered scalable backend  
-✅ **Interactive UI** - Intuitive Streamlit interface  
-✅ **Containerized** - Docker-based deployment  
+✅ **Knowledge Graph Integration** - Neo4j for healthcare data relationships
+✅ **RESTful API** - FastAPI-powered scalable backend
+✅ **Interactive UI** - Intuitive Streamlit interface
+✅ **Containerized** - Docker-based deployment
 ✅ **Multi-Model Support** - Configurable OpenAI models
 
 <a name="architecture"></a>
@@ -147,7 +150,10 @@ Relationships between nodes contain additional contextual information:
 - **Docker**: Containerization platform
 - **OpenAI GPT-3.5**: Language model for natural language understanding
 
-## Acknowledgments
+## 🤝 Contributing
 
-This project builds upon the excellent foundation provided by Real Python's LLM RAG Chatbot [tutorial](https://realpython.com/build-llm-rag-chatbot-with-langchain).
+Contributions are welcome. Please ensure that any pull requests maintain the existing code style and include appropriate tests and documentation updates.
+
+## 🙏 Acknowledgments
 
+This project builds upon the excellent foundation provided by Real Python's LLM RAG Chatbot [tutorial](https://realpython.com/build-llm-rag-chatbot-with-langchain).
diff --git a/chatbot_api/mypy.ini b/chatbot_api/mypy.ini
new file mode 100644
index 0000000..3fa519d
--- /dev/null
+++ b/chatbot_api/mypy.ini
@@ -0,0 +1,41 @@
+[mypy]
+python_version = 3.9
+warn_return_any = True
+warn_unused_configs = True
+disallow_untyped_defs = False
+disallow_incomplete_defs = False
+check_untyped_defs = True
+disallow_untyped_decorators = False
+no_implicit_optional = True
+warn_redundant_casts = True
+warn_unused_ignores = True
+warn_no_return = True
+warn_unreachable = True
+strict_equality = True
+namespace_packages = True
+explicit_package_bases = True
+
+# Ignore missing imports for third-party libraries
+[mypy-langchain.*]
+ignore_missing_imports = True
+
+[mypy-langchain_openai.*]
+ignore_missing_imports = True
+
+[mypy-langchain_community.*]
+ignore_missing_imports = True
+
+[mypy-neo4j.*]
+ignore_missing_imports = True
+
+[mypy-openai.*]
+ignore_missing_imports = True
+
+[mypy-pydantic.*]
+ignore_missing_imports = True
+
+[mypy-fastapi.*]
+ignore_missing_imports = True
+
+[mypy-uvicorn.*]
+ignore_missing_imports = True
diff --git a/chatbot_api/pyproject.toml b/chatbot_api/pyproject.toml
index b317195..a127bb5 100644
--- a/chatbot_api/pyproject.toml
+++ b/chatbot_api/pyproject.toml
@@ -16,4 +16,13 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-dev = ["black", "flake8"]
+dev = [
+    "ruff",
+    "mypy",
+    "pre-commit",
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "pytest-mock>=3.10.0",
+    "httpx>=0.24.0",
+    "pytest-cov>=4.0.0"
+]
diff --git a/chatbot_api/src/agents/hospital_rag_agent.py b/chatbot_api/src/agents/hospital_rag_agent.py
index 6f79bef..9c7211b 100644
--- a/chatbot_api/src/agents/hospital_rag_agent.py
+++ b/chatbot_api/src/agents/hospital_rag_agent.py
@@ -1,4 +1,5 @@
 import os
+from typing import List, Optional
 from langchain_openai import ChatOpenAI
 from langchain.agents import (
     create_openai_functions_agent,
@@ -6,6 +7,7 @@
     AgentExecutor,
 )
 from langchain import hub
+from langchain.schema import BasePromptTemplate
 from chains.hospital_review_chain import reviews_vector_chain
 from chains.hospital_cypher_chain import hospital_cypher_chain
 from tools.wait_times import (
@@ -13,12 +15,12 @@
     get_most_available_hospital,
 )
 
-HOSPITAL_AGENT_MODEL = os.getenv("HOSPITAL_AGENT_MODEL")
+HOSPITAL_AGENT_MODEL: Optional[str] = os.getenv("HOSPITAL_AGENT_MODEL")
 
-hospital_agent_prompt = hub.pull("hwchase17/openai-functions-agent")
+hospital_agent_prompt: BasePromptTemplate = hub.pull("hwchase17/openai-functions-agent")
 
 
-tools = [
+tools: List[Tool] = [
     Tool(
         name="Experiences",
         func=reviews_vector_chain.invoke,
@@ -68,7 +70,7 @@
 ]
 
 
-chat_model = ChatOpenAI(
+chat_model: ChatOpenAI = ChatOpenAI(
     model=HOSPITAL_AGENT_MODEL,
     temperature=0,
 )
@@ -79,7 +81,7 @@
     tools=tools,
 )
 
-hospital_rag_agent_executor = AgentExecutor(
+hospital_rag_agent_executor: AgentExecutor = AgentExecutor(
     agent=hospital_rag_agent,
     tools=tools,
     return_intermediate_steps=True,
diff --git a/chatbot_api/src/chains/hospital_cypher_chain.py b/chatbot_api/src/chains/hospital_cypher_chain.py
index f809b6e..ef4c9df 100644
--- a/chatbot_api/src/chains/hospital_cypher_chain.py
+++ b/chatbot_api/src/chains/hospital_cypher_chain.py
@@ -1,13 +1,14 @@
 import os
+from typing import Optional
 from langchain_community.graphs import Neo4jGraph
 from langchain.chains import GraphCypherQAChain
 from langchain_openai import ChatOpenAI
 from langchain.prompts import PromptTemplate
 
-HOSPITAL_QA_MODEL = os.getenv("HOSPITAL_QA_MODEL")
-HOSPITAL_CYPHER_MODEL = os.getenv("HOSPITAL_CYPHER_MODEL")
+HOSPITAL_QA_MODEL: Optional[str] = os.getenv("HOSPITAL_QA_MODEL")
+HOSPITAL_CYPHER_MODEL: Optional[str] = os.getenv("HOSPITAL_CYPHER_MODEL")
 
-graph = Neo4jGraph(
+graph: Neo4jGraph = Neo4jGraph(
     url=os.getenv("NEO4J_URI"),
     username=os.getenv("NEO4J_USERNAME"),
     password=os.getenv("NEO4J_PASSWORD"),
@@ -16,7 +17,7 @@
 graph.refresh_schema()
 
 
-cypher_generation_template = """
+cypher_generation_template: str = """
 Task:
 Generate Cypher query for a Neo4j graph database.
 
@@ -102,12 +103,12 @@
 {question}
 """
 
-cypher_generation_prompt = PromptTemplate(
+cypher_generation_prompt: PromptTemplate = PromptTemplate(
     input_variables=["schema", "question"], template=cypher_generation_template
 )
 
 
-qa_generation_template = """You are an assistant that takes the results
+qa_generation_template: str = """You are an assistant that takes the results
 from a Neo4j Cypher query and forms a human-readable response. The
 query results section contains the results of a Cypher query that was
 generated based on a users natural language question. The provided
@@ -141,12 +142,12 @@
 Helpful Answer:
 """
 
-qa_generation_prompt = PromptTemplate(
+qa_generation_prompt: PromptTemplate = PromptTemplate(
     input_variables=["context", "question"], template=qa_generation_template
 )
 
 
-hospital_cypher_chain = GraphCypherQAChain.from_llm(
+hospital_cypher_chain: GraphCypherQAChain = GraphCypherQAChain.from_llm(
     cypher_llm=ChatOpenAI(model=HOSPITAL_CYPHER_MODEL, temperature=0),
     qa_llm=ChatOpenAI(model=HOSPITAL_QA_MODEL, temperature=0),
     graph=graph,
diff --git a/chatbot_api/src/chains/hospital_review_chain.py b/chatbot_api/src/chains/hospital_review_chain.py
index a182e58..7aa4077 100644
--- a/chatbot_api/src/chains/hospital_review_chain.py
+++ b/chatbot_api/src/chains/hospital_review_chain.py
@@ -1,4 +1,5 @@
 import os
+from typing import Optional, List, Union
 from langchain.vectorstores.neo4j_vector import Neo4jVector
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain.chains import RetrievalQA
@@ -9,9 +10,9 @@
     ChatPromptTemplate,
 )
 
-HOSPITAL_QA_MODEL = os.getenv("HOSPITAL_QA_MODEL")
+HOSPITAL_QA_MODEL: Optional[str] = os.getenv("HOSPITAL_QA_MODEL")
 
-neo4j_vector_index = Neo4jVector.from_existing_graph(
+neo4j_vector_index: Neo4jVector = Neo4jVector.from_existing_graph(
     embedding=OpenAIEmbeddings(),
     url=os.getenv("NEO4J_URI"),
     username=os.getenv("NEO4J_USERNAME"),
@@ -27,7 +28,7 @@
     embedding_node_property="embedding",
 )
 
-review_template = """Your job is to use patient
+review_template: str = """Your job is to use patient
 reviews to answer questions about their experience at a hospital. Use
 the following context to answer questions. Be as detailed as possible, but
 don't make up any information that's not from the context. If you don't know
@@ -35,20 +36,23 @@
 {context}
 """
 
-review_system_prompt = SystemMessagePromptTemplate(
+review_system_prompt: SystemMessagePromptTemplate = SystemMessagePromptTemplate(
     prompt=PromptTemplate(input_variables=["context"], template=review_template)
 )
 
-review_human_prompt = HumanMessagePromptTemplate(
+review_human_prompt: HumanMessagePromptTemplate = HumanMessagePromptTemplate(
     prompt=PromptTemplate(input_variables=["question"], template="{question}")
 )
-messages = [review_system_prompt, review_human_prompt]
+messages: List[Union[SystemMessagePromptTemplate, HumanMessagePromptTemplate]] = [
+    review_system_prompt,
+    review_human_prompt,
+]
 
-review_prompt = ChatPromptTemplate(
+review_prompt: ChatPromptTemplate = ChatPromptTemplate(
     input_variables=["context", "question"], messages=messages
 )
 
-reviews_vector_chain = RetrievalQA.from_chain_type(
+reviews_vector_chain: RetrievalQA = RetrievalQA.from_chain_type(
     llm=ChatOpenAI(model=HOSPITAL_QA_MODEL, temperature=0),
     chain_type="stuff",
     retriever=neo4j_vector_index.as_retriever(k=12),
diff --git a/chatbot_api/src/main.py b/chatbot_api/src/main.py
index 0e4cf2a..8c8697a 100644
--- a/chatbot_api/src/main.py
+++ b/chatbot_api/src/main.py
@@ -1,17 +1,18 @@
+from typing import Dict, Any
 from fastapi import FastAPI
 from agents.hospital_rag_agent import hospital_rag_agent_executor
 from models.hospital_rag_query import HospitalQueryInput, HospitalQueryOutput
 from utils.async_utils import async_retry
 
 
-app = FastAPI(
+app: FastAPI = FastAPI(
     title="Hospital Chatbot",
     description="Endpoints for a hospital system graph RAG chatbot",
 )
 
 
 @async_retry(max_retries=10, delay=1)
-async def invoke_agent_with_retry(query: str):
+async def invoke_agent_with_retry(query: str) -> Dict[str, Any]:
     """Retry the agent if a tool fails to run.
 
     This can help when there are intermittent connection issues
@@ -21,13 +22,13 @@ async def invoke_agent_with_retry(query: str):
 
 
 @app.get("/")
-async def get_status():
+async def get_status() -> Dict[str, str]:
     return {"status": "running"}
 
 
 @app.post("/hospital-rag-agent")
 async def query_hospital_agent(query: HospitalQueryInput) -> HospitalQueryOutput:
-    query_response = await invoke_agent_with_retry(query.text)
+    query_response: Dict[str, Any] = await invoke_agent_with_retry(query.text)
     query_response["intermediate_steps"] = [
         str(s) for s in query_response["intermediate_steps"]
     ]
diff --git a/chatbot_api/src/tools/wait_times.py b/chatbot_api/src/tools/wait_times.py
index 24d2ad8..d0e5d26 100644
--- a/chatbot_api/src/tools/wait_times.py
+++ b/chatbot_api/src/tools/wait_times.py
@@ -52,9 +52,7 @@ def get_most_available_hospital(_: Any) -> dict[str, float]:
     """Find the hospital with the shortest wait time."""
     current_hospitals = _get_current_hospitals()
 
-    current_wait_times = [
-        _get_current_wait_time_minutes(h) for h in current_hospitals
-    ]
+    current_wait_times = [_get_current_wait_time_minutes(h) for h in current_hospitals]
 
     best_time_idx = np.argmin(current_wait_times)
     best_hospital = current_hospitals[best_time_idx]
diff --git a/chatbot_frontend/pyproject.toml b/chatbot_frontend/pyproject.toml
index dacd56d..9b4d90a 100644
--- a/chatbot_frontend/pyproject.toml
+++ b/chatbot_frontend/pyproject.toml
@@ -7,4 +7,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-dev = ["black", "flake8"]
+dev = [
+    "ruff",
+    "mypy",
+    "pre-commit",
+    "pytest>=7.0.0",
+    "pytest-mock>=3.10.0",
+    "httpx>=0.24.0",
+    "pytest-cov>=4.0.0"
+]
diff --git a/chatbot_frontend/src/main.py b/chatbot_frontend/src/main.py
index becd16b..92c709e 100644
--- a/chatbot_frontend/src/main.py
+++ b/chatbot_frontend/src/main.py
@@ -26,8 +26,7 @@
     )
     st.markdown("- What is the average duration in days for closed emergency visits?")
     st.markdown(
-        "- What are patients saying about the nursing staff at "
-        "Castaneda-Hardy?"
+        "- What are patients saying about the nursing staff at Castaneda-Hardy?"
     )
     st.markdown("- What was the total billing amount charged to each payer for 2023?")
     st.markdown("- What is the average billing amount for medicaid visits?")
@@ -101,4 +100,4 @@
             "output": output_text,
             "explanation": explanation,
         }
-    )
\ No newline at end of file
+    )
diff --git a/docker-compose.ci.yml b/docker-compose.ci.yml
new file mode 100644
index 0000000..3080783
--- /dev/null
+++ b/docker-compose.ci.yml
@@ -0,0 +1,26 @@
+services:
+  hospital_neo4j_etl:
+    environment:
+      # Dummy environment variables for CI build validation
+      - NEO4J_URI=neo4j://localhost:7687
+      - NEO4J_USERNAME=neo4j
+      - NEO4J_PASSWORD=password
+    env_file: null  # Completely disable env_file
+
+  chatbot_api:
+    environment:
+      # Dummy environment variables for CI build validation
+      - OPENAI_API_KEY=sk-dummy-key-for-ci-build-only
+      - NEO4J_URI=neo4j://localhost:7687
+      - NEO4J_USERNAME=neo4j
+      - NEO4J_PASSWORD=password
+      - HOSPITAL_AGENT_MODEL=gpt-3.5-turbo
+      - HOSPITAL_CYPHER_MODEL=gpt-3.5-turbo
+      - HOSPITAL_QA_MODEL=gpt-3.5-turbo
+    env_file: null  # Completely disable env_file
+
+  chatbot_frontend:
+    environment:
+      # Dummy environment variables for CI build validation
+      - CHATBOT_URL=http://localhost:8000/hospital-rag-agent
+    env_file: null  # Completely disable env_file
diff --git a/docker-compose.yml b/docker-compose.yml
index c9ec18e..3a02916 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,4 +1,3 @@
-version: '3'
 
 services:
   hospital_neo4j_etl:
diff --git a/hospital_neo4j_etl/Dockerfile b/hospital_neo4j_etl/Dockerfile
index 71ff24f..69783f5 100644
--- a/hospital_neo4j_etl/Dockerfile
+++ b/hospital_neo4j_etl/Dockerfile
@@ -7,4 +7,4 @@ COPY ./src/ /app
 COPY ./pyproject.toml /code/pyproject.toml
 RUN pip install /code/.
 
-CMD ["sh", "entrypoint.sh"]
\ No newline at end of file
+CMD ["sh", "entrypoint.sh"]
diff --git a/hospital_neo4j_etl/pyproject.toml b/hospital_neo4j_etl/pyproject.toml
index 35c7467..ef30ae7 100644
--- a/hospital_neo4j_etl/pyproject.toml
+++ b/hospital_neo4j_etl/pyproject.toml
@@ -7,4 +7,11 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-dev = ["black", "flake8"]
\ No newline at end of file
+dev = [
+    "ruff",
+    "mypy",
+    "pre-commit",
+    "pytest>=7.0.0",
+    "pytest-mock>=3.10.0",
+    "pytest-cov>=4.0.0"
+]
diff --git a/hospital_neo4j_etl/src/hospital_bulk_csv_write.py b/hospital_neo4j_etl/src/hospital_bulk_csv_write.py
index ccbac75..78edd25 100644
--- a/hospital_neo4j_etl/src/hospital_bulk_csv_write.py
+++ b/hospital_neo4j_etl/src/hospital_bulk_csv_write.py
@@ -38,9 +38,7 @@ def load_hospital_graph_from_csv() -> None:
     """Load structured hospital CSV data following
     a specific ontology into Neo4j"""
 
-    driver = GraphDatabase.driver(
-        NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)
-    )
+    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
 
     LOGGER.info("Setting uniqueness constraints on nodes")
     with driver.session(database="neo4j") as session:
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..798b4e2
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,22 @@
+[tool:pytest]
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*
+addopts =
+    --verbose
+    --tb=short
+    --cov-report=term-missing
+    --cov-report=html
+    --cov-fail-under=75
+    --cov-config=.coveragerc
+    --ignore=tests/performance
+asyncio_mode = auto
+markers =
+    slow: marks tests as slow (deselect with '-m "not slow"')
+env =
+    ENVIRONMENT = test
+    NEO4J_URI = neo4j://localhost:7687
+    NEO4J_USERNAME = test
+    NEO4J_PASSWORD = test
+    OPENAI_API_KEY = test-key
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/async_agent_requests.py b/tests/async_agent_requests.py
deleted file mode 100644
index 02a9454..0000000
--- a/tests/async_agent_requests.py
+++ /dev/null
@@ -1,45 +0,0 @@
-import asyncio
-import time
-import httpx
-
-CHATBOT_URL = "http://localhost:8000/hospital-rag-agent"
-
-
-async def make_async_post(url, data):
-    timeout = httpx.Timeout(timeout=120)
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url, json=data, timeout=timeout)
-        return response
-
-
-async def make_bulk_requests(url, data):
-    tasks = [make_async_post(url, payload) for payload in data]
-    responses = await asyncio.gather(*tasks)
-    outputs = [r.json()["output"] for r in responses]
-    return outputs
-
-
-questions = [
-   "What is the current wait time at Wallace-Hamilton hospital?",
-   "Which hospital has the shortest wait time?",
-   "At which hospitals are patients complaining about billing and insurance issues?",
-   "What is the average duration in days for emergency visits?",
-   "What are patients saying about the nursing staff at Castaneda-Hardy?",
-   "What was the total billing amount charged to each payer for 2023?",
-   "What is the average billing amount for medicaid visits?",
-   "How many patients has Dr. Ryan Brown treated?",
-   "Which physician has the lowest average visit duration in days?",
-   "How many visits are open and what is their average duration in days?",
-   "Have any patients complained about noise?",
-   "How much was billed for patient 789's stay?",
-   "Which physician has billed the most to cigna?",
-   "Which state had the largest percent increase in medicaid visits from 2022 to 2023?",
-]
-
-request_bodies = [{"text": q} for q in questions]
-
-start_time = time.perf_counter()
-outputs = asyncio.run(make_bulk_requests(CHATBOT_URL, request_bodies))
-end_time = time.perf_counter()
-
-print(f"Run time: {end_time - start_time} seconds")
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..fb4a7fb
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,53 @@
+import pytest
+import os
+from unittest.mock import Mock, AsyncMock
+from typing import Generator
+
+
+@pytest.fixture(scope="session")
+def test_env() -> Generator[None, None, None]:
+    os.environ["ENVIRONMENT"] = "test"
+    os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
+    os.environ["NEO4J_USERNAME"] = "test"
+    os.environ["NEO4J_PASSWORD"] = "test"
+    os.environ["OPENAI_API_KEY"] = "test-key"
+    os.environ["HOSPITAL_AGENT_MODEL"] = "gpt-3.5-turbo"
+    os.environ["HOSPITAL_CYPHER_MODEL"] = "gpt-3.5-turbo"
+    os.environ["HOSPITAL_QA_MODEL"] = "gpt-3.5-turbo"
+    yield
+
+
+@pytest.fixture
+def mock_neo4j_graph():
+    mock_graph = Mock()
+    mock_graph.query.return_value = []
+    mock_graph.refresh_schema.return_value = None
+    return mock_graph
+
+
+@pytest.fixture
+def mock_openai_client():
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create.return_value = Mock(
+        choices=[Mock(message=Mock(content="Test response"))]
+    )
+    return mock_client
+
+
+@pytest.fixture
+def sample_hospital_data():
+    return [{"hospital_name": "Test Hospital"}, {"hospital_name": "Another Hospital"}]
+
+
+@pytest.fixture
+def sample_query_input():
+    return {"text": "What is the current wait time at Test Hospital?"}
+
+
+@pytest.fixture
+def sample_agent_response():
+    return {
+        "input": "Test query",
+        "output": "Test response",
+        "intermediate_steps": ["Step 1", "Step 2"],
+    }
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/integration/test_api.py b/tests/integration/test_api.py
new file mode 100644
index 0000000..abdf900
--- /dev/null
+++ b/tests/integration/test_api.py
@@ -0,0 +1,38 @@
+import pytest
+
+from chatbot_api.src.utils.async_utils import async_retry
+from chatbot_api.src.models.hospital_rag_query import (
+    HospitalQueryInput,
+    HospitalQueryOutput,
+)
+
+
+class TestWorkflowIntegration:
+    @pytest.mark.asyncio
+    async def test_async_retry_with_business_logic(self):
+        """Test async_retry works with our business functions."""
+        call_count = 0
+
+        @async_retry(max_retries=3, delay=0.01)
+        async def unreliable_data_processor(input_text):
+            nonlocal call_count
+            call_count += 1
+
+            if call_count < 2:
+                raise Exception("Processing failed")
+
+            # Our business logic: process input and create output
+            query_input = HospitalQueryInput(text=input_text)
+            query_output = HospitalQueryOutput(
+                input=query_input.text,
+                output=f"Processed: {query_input.text}",
+                intermediate_steps=[f"Attempt {call_count}", "Success"],
+            )
+            return query_output
+
+        result = await unreliable_data_processor("Test query")
+
+        assert result.input == "Test query"
+        assert "Processed: Test query" == result.output
+        assert len(result.intermediate_steps) == 2
+        assert call_count == 2  # Failed once, succeeded on second try
diff --git a/tests/performance/__init__.py b/tests/performance/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/performance/test_agent_performance.py b/tests/performance/test_agent_performance.py
new file mode 100644
index 0000000..89f1bc9
--- /dev/null
+++ b/tests/performance/test_agent_performance.py
@@ -0,0 +1,113 @@
+import pytest
+import asyncio
+import time
+import httpx
+
+
+class TestAgentPerformance:
+    """Performance tests for the hospital RAG agent."""
+
+    @pytest.fixture
+    def test_questions(self):
+        return [
+            "What is the current wait time at Wallace-Hamilton hospital?",
+            "Which hospital has the shortest wait time?",
+            "At which hospitals are patients complaining about billing and insurance issues?",
+            "What is the average duration in days for emergency visits?",
+            "What are patients saying about the nursing staff at Castaneda-Hardy?",
+            "What was the total billing amount charged to each payer for 2023?",
+            "What is the average billing amount for medicaid visits?",
+            "How many patients has Dr. Ryan Brown treated?",
+            "Which physician has the lowest average visit duration in days?",
+            "How many visits are open and what is their average duration in days?",
+            "Have any patients complained about noise?",
+            "How much was billed for patient 789's stay?",
+            "Which physician has billed the most to cigna?",
+            "Which state had the largest percent increase in medicaid visits from 2022 to 2023?",
+        ]
+
+    @pytest.fixture
+    def chatbot_url(self):
+        return "http://localhost:8000/hospital-rag-agent"
+
+    @pytest.mark.asyncio
+    @pytest.mark.slow
+    async def test_async_performance_benchmark(self, test_questions, chatbot_url):
+        """Test async performance of multiple concurrent requests."""
+
+        async def make_request(question: str) -> httpx.Response:
+            timeout = httpx.Timeout(timeout=120)
+            async with httpx.AsyncClient() as client:
+                return await client.post(
+                    chatbot_url, json={"text": question}, timeout=timeout
+                )
+
+        start_time = time.perf_counter()
+        tasks = [make_request(q) for q in test_questions]
+        responses = await asyncio.gather(*tasks, return_exceptions=True)
+        end_time = time.perf_counter()
+
+        total_time = end_time - start_time
+        successful_responses = [
+            r
+            for r in responses
+            if isinstance(r, httpx.Response) and r.status_code == 200
+        ]
+
+        assert len(successful_responses) > 0, "No successful responses received"
+        assert total_time < 300, f"Async requests took too long: {total_time}s"
+
+        print(f"Async performance: {len(test_questions)} requests in {total_time:.2f}s")
+        print(
+            f"Successful responses: {len(successful_responses)}/{len(test_questions)}"
+        )
+
+    @pytest.mark.slow
+    def test_sync_performance_benchmark(self, test_questions, chatbot_url):
+        """Test synchronous performance for comparison."""
+        import requests
+
+        start_time = time.perf_counter()
+        responses = []
+        for question in test_questions:
+            try:
+                response = requests.post(
+                    chatbot_url, json={"text": question}, timeout=120
+                )
+                responses.append(response)
+            except Exception as e:
+                responses.append(e)
+        end_time = time.perf_counter()
+
+        total_time = end_time - start_time
+        successful_responses = [
+            r for r in responses if hasattr(r, "status_code") and r.status_code == 200
+        ]
+
+        assert len(successful_responses) > 0, "No successful responses received"
+
+        print(f"Sync performance: {len(test_questions)} requests in {total_time:.2f}s")
+        print(
+            f"Successful responses: {len(successful_responses)}/{len(test_questions)}"
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.slow
+    async def test_single_request_response_time(self, chatbot_url):
+        """Test response time for a single request."""
+        question = "What is the current wait time at Wallace-Hamilton hospital?"
+
+        start_time = time.perf_counter()
+        timeout = httpx.Timeout(timeout=120)
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                chatbot_url, json={"text": question}, timeout=timeout
+            )
+        end_time = time.perf_counter()
+
+        response_time = end_time - start_time
+
+        assert response.status_code == 200
+        assert response_time < 30, f"Single request took too long: {response_time}s"
+
+        print(f"Single request response time: {response_time:.2f}s")
diff --git a/tests/sync_agent_requests.py b/tests/sync_agent_requests.py
deleted file mode 100644
index 02aae7c..0000000
--- a/tests/sync_agent_requests.py
+++ /dev/null
@@ -1,29 +0,0 @@
-import time
-import requests
-
-CHATBOT_URL = "http://localhost:8000/hospital-rag-agent"
-
-questions = [
-   "What is the current wait time at Wallace-Hamilton hospital?",
-   "Which hospital has the shortest wait time?",
-   "At which hospitals are patients complaining about billing and insurance issues?",
-   "What is the average duration in days for emergency visits?",
-   "What are patients saying about the nursing staff at Castaneda-Hardy?",
-   "What was the total billing amount charged to each payer for 2023?",
-   "What is the average billing amount for medicaid visits?",
-   "How many patients has Dr. Ryan Brown treated?",
-   "Which physician has the lowest average visit duration in days?",
-   "How many visits are open and what is their average duration in days?",
-   "Have any patients complained about noise?",
-   "How much was billed for patient 789's stay?",
-   "Which physician has billed the most to cigna?",
-   "Which state had the largest percent increase in medicaid visits from 2022 to 2023?",
-]
-
-request_bodies = [{"text": q} for q in questions]
-
-start_time = time.perf_counter()
-outputs = [requests.post(CHATBOT_URL, json=data) for data in request_bodies]
-end_time = time.perf_counter()
-
-print(f"Run time: {end_time - start_time} seconds")
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unit/chatbot_api/__init__.py b/tests/unit/chatbot_api/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unit/chatbot_api/test_async_utils.py b/tests/unit/chatbot_api/test_async_utils.py
new file mode 100644
index 0000000..8935aa6
--- /dev/null
+++ b/tests/unit/chatbot_api/test_async_utils.py
@@ -0,0 +1,72 @@
+import pytest
+from unittest.mock import patch
+
+from chatbot_api.src.utils.async_utils import async_retry
+
+
+class TestAsyncRetry:
+    @pytest.mark.asyncio
+    async def test_success_on_first_attempt(self):
+        @async_retry(max_retries=3, delay=0.1)
+        async def successful_function():
+            return "success"
+
+        result = await successful_function()
+        assert result == "success"
+
+    @pytest.mark.asyncio
+    async def test_success_after_retries(self):
+        call_count = 0
+
+        @async_retry(max_retries=3, delay=0.1)
+        async def function_with_retries():
+            nonlocal call_count
+            call_count += 1
+            if call_count < 3:
+                raise Exception("Temporary failure")
+            return "success"
+
+        result = await function_with_retries()
+        assert result == "success"
+        assert call_count == 3
+
+    @pytest.mark.asyncio
+    async def test_failure_after_max_retries(self):
+        call_count = 0
+
+        @async_retry(max_retries=2, delay=0.1)
+        async def always_failing_function():
+            nonlocal call_count
+            call_count += 1
+            raise Exception("Always fails")
+
+        with pytest.raises(ValueError, match="Failed after 2 attempts"):
+            await always_failing_function()
+
+        assert call_count == 2
+
+    @pytest.mark.asyncio
+    async def test_custom_delay(self):
+        call_count = 0
+
+        @async_retry(max_retries=2, delay=0.01)
+        async def function_with_custom_delay():
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise Exception("First failure")
+            return "success"
+
+        with patch("asyncio.sleep") as mock_sleep:
+            result = await function_with_custom_delay()
+            assert result == "success"
+            mock_sleep.assert_called_once_with(0.01)
+
+    @pytest.mark.asyncio
+    async def test_preserves_function_args_and_kwargs(self):
+        @async_retry(max_retries=1, delay=0.1)
+        async def function_with_args(arg1, arg2, kwarg1=None):
+            return f"{arg1}-{arg2}-{kwarg1}"
+
+        result = await function_with_args("a", "b", kwarg1="c")
+        assert result == "a-b-c"
diff --git a/tests/unit/chatbot_api/test_wait_times.py b/tests/unit/chatbot_api/test_wait_times.py
new file mode 100644
index 0000000..636994c
--- /dev/null
+++ b/tests/unit/chatbot_api/test_wait_times.py
@@ -0,0 +1,133 @@
+import pytest
+from unittest.mock import patch, Mock
+
+from chatbot_api.src.tools.wait_times import (
+    _get_current_hospitals,
+    _get_current_wait_time_minutes,
+    get_current_wait_times,
+    get_most_available_hospital,
+)
+
+
+class TestGetCurrentHospitals:
+    @patch("chatbot_api.src.tools.wait_times.Neo4jGraph")
+    def test_returns_hospital_names(self, mock_neo4j_graph):
+        mock_graph_instance = Mock()
+        mock_graph_instance.query.return_value = [
+            {"hospital_name": "Hospital A"},
+            {"hospital_name": "Hospital B"},
+        ]
+        mock_neo4j_graph.return_value = mock_graph_instance
+
+        result = _get_current_hospitals()
+
+        assert result == ["hospital a", "hospital b"]
+        mock_graph_instance.query.assert_called_once()
+
+    @patch("chatbot_api.src.tools.wait_times.Neo4jGraph")
+    def test_handles_empty_result(self, mock_neo4j_graph):
+        mock_graph_instance = Mock()
+        mock_graph_instance.query.return_value = []
+        mock_neo4j_graph.return_value = mock_graph_instance
+
+        result = _get_current_hospitals()
+
+        assert result == []
+
+
+class TestGetCurrentWaitTimeMinutes:
+    @patch("chatbot_api.src.tools.wait_times._get_current_hospitals")
+    @patch("tools.wait_times.np.random.randint")
+    def test_returns_wait_time_for_valid_hospital(
+        self, mock_randint, mock_get_hospitals
+    ):
+        mock_get_hospitals.return_value = ["test hospital", "another hospital"]
+        mock_randint.return_value = 45
+
+        result = _get_current_wait_time_minutes("Test Hospital")
+
+        assert result == 45
+        mock_randint.assert_called_once_with(low=0, high=600)
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_hospitals")
+    def test_returns_negative_one_for_invalid_hospital(self, mock_get_hospitals):
+        mock_get_hospitals.return_value = ["test hospital", "another hospital"]
+
+        result = _get_current_wait_time_minutes("Nonexistent Hospital")
+
+        assert result == -1
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_hospitals")
+    @patch("tools.wait_times.np.random.randint")
+    def test_case_insensitive_hospital_matching(self, mock_randint, mock_get_hospitals):
+        mock_get_hospitals.return_value = ["test hospital"]
+        mock_randint.return_value = 30
+
+        result = _get_current_wait_time_minutes("TEST HOSPITAL")
+
+        assert result == 30
+
+
+class TestGetCurrentWaitTimes:
+    @patch("chatbot_api.src.tools.wait_times._get_current_wait_time_minutes")
+    def test_formats_wait_time_minutes_only(self, mock_get_wait_time):
+        mock_get_wait_time.return_value = 45
+
+        result = get_current_wait_times("Test Hospital")
+
+        assert result == "45 minutes"
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_wait_time_minutes")
+    def test_formats_wait_time_hours_and_minutes(self, mock_get_wait_time):
+        mock_get_wait_time.return_value = 125  # 2 hours 5 minutes
+
+        result = get_current_wait_times("Test Hospital")
+
+        assert result == "2 hours 5 minutes"
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_wait_time_minutes")
+    def test_formats_wait_time_exact_hours(self, mock_get_wait_time):
+        mock_get_wait_time.return_value = 120  # 2 hours 0 minutes
+
+        result = get_current_wait_times("Test Hospital")
+
+        assert result == "2 hours 0 minutes"
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_wait_time_minutes")
+    def test_handles_nonexistent_hospital(self, mock_get_wait_time):
+        mock_get_wait_time.return_value = -1
+
+        result = get_current_wait_times("Nonexistent Hospital")
+
+        assert result == "Hospital 'Nonexistent Hospital' does not exist."
+
+
+class TestGetMostAvailableHospital:
+    @patch("chatbot_api.src.tools.wait_times._get_current_hospitals")
+    @patch("chatbot_api.src.tools.wait_times._get_current_wait_time_minutes")
+    def test_returns_hospital_with_shortest_wait(
+        self, mock_get_wait_time, mock_get_hospitals
+    ):
+        mock_get_hospitals.return_value = ["hospital a", "hospital b", "hospital c"]
+        mock_get_wait_time.side_effect = [60, 30, 45]  # hospital b has shortest wait
+
+        result = get_most_available_hospital(None)
+
+        assert result == {"hospital b": 30}
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_hospitals")
+    @patch("chatbot_api.src.tools.wait_times._get_current_wait_time_minutes")
+    def test_handles_single_hospital(self, mock_get_wait_time, mock_get_hospitals):
+        mock_get_hospitals.return_value = ["only hospital"]
+        mock_get_wait_time.return_value = 90
+
+        result = get_most_available_hospital(None)
+
+        assert result == {"only hospital": 90}
+
+    @patch("chatbot_api.src.tools.wait_times._get_current_hospitals")
+    def test_handles_no_hospitals(self, mock_get_hospitals):
+        mock_get_hospitals.return_value = []
+
+        with pytest.raises(ValueError):
+            get_most_available_hospital(None)
diff --git a/tests/unit/chatbot_frontend/__init__.py b/tests/unit/chatbot_frontend/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unit/hospital_neo4j_etl/__init__.py b/tests/unit/hospital_neo4j_etl/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unit/hospital_neo4j_etl/test_hospital_bulk_csv_write.py b/tests/unit/hospital_neo4j_etl/test_hospital_bulk_csv_write.py
new file mode 100644
index 0000000..c63b4d5
--- /dev/null
+++ b/tests/unit/hospital_neo4j_etl/test_hospital_bulk_csv_write.py
@@ -0,0 +1,31 @@
+from unittest.mock import Mock
+
+from hospital_neo4j_etl.src.hospital_bulk_csv_write import (
+    _set_uniqueness_constraints,
+    NODES,
+)
+
+
+class TestSetUniquenessConstraints:
+    def test_generates_correct_constraint_query(self):
+        """Test our constraint query generation logic."""
+        mock_tx = Mock()
+
+        _set_uniqueness_constraints(mock_tx, "Hospital")
+
+        expected_query = "CREATE CONSTRAINT IF NOT EXISTS FOR (n:Hospital)\n        REQUIRE n.id IS UNIQUE;"
+        mock_tx.run.assert_called_once_with(expected_query, {})
+
+    def test_works_with_all_defined_node_types(self):
+        """Test constraint generation works for all our node types."""
+        mock_tx = Mock()
+
+        for node in NODES:
+            mock_tx.reset_mock()
+            _set_uniqueness_constraints(mock_tx, node)
+
+            call_args = mock_tx.run.call_args[0]
+            query = call_args[0]
+
+            assert f"FOR (n:{node})" in query
+            assert "REQUIRE n.id IS UNIQUE" in query