feat: set up integration tests

Author: Eyo Chen

Makefile

@@ -1,2 +1,5 @@
 generate:
-	docker-compose run --rm -T protoc
+	docker-compose run --rm -T protoc
+
+test:
+	./tools/tests/run_tests.sh

docker-compose.yml

@@ -1,6 +1,13 @@
 version: '3.8'
 
 services:
+  mongodb-test:
+    image: mongo:latest
+    container_name: test_mongodb
+    ports:
+      - "27017:27017"
+    environment:
+      MONGO_INITDB_DATABASE: test_stock_db
   protoc:
     build:
       context: ./tools/protoc

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "markupsafe==3.0.2",
     "polygon-api-client==1.14.5",
     "protobuf==5.29.4",
+    "pymongo>=4.13.0",
     "python-dateutil==2.9.0.post0",
     "python-dotenv==1.1.0",
     "pyyaml==6.0.2",
@@ -25,3 +26,8 @@ dependencies = [
     "validate-email==1.3",
     "websockets==14.2",
 ]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3.5",
+]

pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+python_files = test_*.py
+python_functions = test_*
+addopts = -v --tb=short
+testpaths = src/tests

src/tests/tools/wait_for_mongo.py

@@ -0,0 +1,26 @@
+import time
+from pymongo import MongoClient
+from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
+
+
+def wait_for_mongo(host="localhost", port=27017, timeout=30, db_name="test_stock_db"):
+    """Wait for MongoDB to be ready by attempting to connect and ping."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            client = MongoClient(
+                f"mongodb://{host}:{port}", serverSelectionTimeoutMS=1000
+            )
+            client.admin.command("ping")  # Ping the server to check if it's ready
+            client.drop_database(db_name)  # Ensure clean state
+            client.close()
+            print("MongoDB is ready and test database cleared!")
+            return True
+        except (ConnectionFailure, ServerSelectionTimeoutError):
+            print("MongoDB not ready, retrying...")
+            time.sleep(1)
+    raise TimeoutError("MongoDB did not become ready within the timeout period")
+
+
+if __name__ == "__main__":
+    wait_for_mongo()

tools/tests/run_tests.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+
+# Define color codes
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+BLUE='\033[0;34m'
+YELLOW='\033[0;33m'
+NC='\033[0m'
+
+# Function to print a line of repeated emojis
+print_emoji_line() {
+    local emoji=$1
+    local color=$2
+    local cols=$(tput cols) # Get terminal width
+    local line=""
+
+    # Calculate how many emojis fit in the terminal width (approximate, as emojis vary in width)
+    for ((i=0; i<cols/2; i++)); do
+        line="${line}${emoji}"
+    done
+    echo -e "${color}${line}"
+}
+
+# Function to clean up resources
+cleanup() {
+    echo -e "${BLUE}"
+    print_emoji_line "<=" "${BLUE}"
+    echo -e "<= Cleaning up resources..."
+    print_emoji_line "<=" "${BLUE}"
+    docker-compose down -v mongodb-test
+    rm -rf __pycache__ tests/__pycache__
+}
+
+# Set trap to call cleanup on exit (success or failure)
+trap cleanup EXIT
+
+echo -e "${YELLOW}"
+print_emoji_line "=>" "${YELLOW}"
+echo -e "=> Starting test environment..."
+print_emoji_line "=>" "${YELLOW}"
+docker-compose up -d mongodb-test
+uv run src/tests/tools/wait_for_mongo.py
+PYTHONPATH=./src python -m pytest src/tests/
+if [ $? -eq 0 ]; then
+    echo -e "${GREEN}🎉 Tests passed successfully!"
+else
+    echo -e "${RED}❌ Tests failed!"
+fi