From 1c9790fd7bd93adebd1716e289893faf99bea909 Mon Sep 17 00:00:00 2001 From: Mohammad Amin Date: Thu, 28 Dec 2023 17:26:19 +0330 Subject: [PATCH] feat: Added rabbitmq and .env variable reading! --- celery_app/server.py | 11 +++- celery_app/tasks.py | 10 ++-- docker-compose.test.yml | 15 +++++ tests/unit/test_load_credentials.py | 41 ++++++++++++++ .../test_load_retriever_hyperparameters.py | 2 +- utils/__init__.py | 0 utils/credentials.py | 55 +++++++++++++++++++ worker.py | 12 ++-- 8 files changed, 134 insertions(+), 12 deletions(-) create mode 100644 tests/unit/test_load_credentials.py create mode 100644 utils/__init__.py create mode 100644 utils/credentials.py diff --git a/celery_app/server.py b/celery_app/server.py index c1c44d4..e9e2743 100644 --- a/celery_app/server.py +++ b/celery_app/server.py @@ -1,5 +1,12 @@ from celery import Celery -# TODO: read from .env -app = Celery("celery_app/tasks", broker="pyamqp://root:pass@localhost//") +from utils.credentials import load_rabbitmq_credentials + +rabbit_creds = load_rabbitmq_credentials() +user = rabbit_creds['user'] +password = rabbit_creds['password'] +host = rabbit_creds['host'] +port = rabbit_creds['port'] + +app = Celery("celery_app/tasks", broker=f"pyamqp://{user}:{password}@{host}:{port}//") app.autodiscover_tasks(["celery_app"]) diff --git a/celery_app/tasks.py b/celery_app/tasks.py index e2025af..b4fea92 100644 --- a/celery_app/tasks.py +++ b/celery_app/tasks.py @@ -1,15 +1,17 @@ from celery_app.job_send import job_send from celery_app.server import app +from utils.credentials import load_rabbitmq_credentials # TODO: Write tasks that match our requirements @app.task def add(x, y): - broker_url = "localhost" - port = 5672 - username = "root" - password = "pass" + rabbit_creds = load_rabbitmq_credentials() + username = rabbit_creds['user'] + password = rabbit_creds['password'] + broker_url = rabbit_creds['host'] + port = rabbit_creds['port'] res = x + y job_send(broker_url, port, username, password, res) diff --git a/docker-compose.test.yml b/docker-compose.test.yml index 97fbcea..4e046d9 100644 --- a/docker-compose.test.yml +++ b/docker-compose.test.yml @@ -22,6 +22,10 @@ services: - POSTGRES_USER=root - POSTGRES_PASS=pass - POSTGRES_PORT=5432 + - RABBIT_HOST=rabbitmq + - RABBIT_PORT=5672 + - RABBIT_USER=root + - RABBIT_PASSWORD=pass - CHUNK_SIZE=512 - EMBEDDING_DIM=1024 - K1_RETRIEVER_SEARCH=20 @@ -69,3 +73,14 @@ services: interval: 10s timeout: 5s retries: 5 + rabbitmq: + image: "rabbitmq:3-management-alpine" + environment: + - RABBITMQ_DEFAULT_USER=root + - RABBITMQ_DEFAULT_PASS=pass + healthcheck: + test: rabbitmq-diagnostics -q ping + interval: 30s + timeout: 30s + retries: 2 + start_period: 40s diff --git a/tests/unit/test_load_credentials.py b/tests/unit/test_load_credentials.py new file mode 100644 index 0000000..701b5e6 --- /dev/null +++ b/tests/unit/test_load_credentials.py @@ -0,0 +1,41 @@ +import unittest + +from utils.credentials import load_postgres_credentials, load_rabbitmq_credentials + + +class TestCredentialLoadings(unittest.TestCase): + def test_postgresql_envs_check_type(self): + postgres_creds = load_postgres_credentials() + + self.assertIsInstance(postgres_creds, dict) + + def test_postgresql_envs_values(self): + postgres_creds = load_postgres_credentials() + + self.assertNotEqual(postgres_creds["user"], None) + self.assertNotEqual(postgres_creds["password"], None) + self.assertNotEqual(postgres_creds["host"], None) + self.assertNotEqual(postgres_creds["port"], None) + + self.assertIsInstance(postgres_creds["user"], str) + self.assertIsInstance(postgres_creds["password"], str) + self.assertIsInstance(postgres_creds["host"], str) + self.assertIsInstance(postgres_creds["port"], str) + + def test_rabbitmq_envs_check_type(self): + rabbitmq_creds = load_rabbitmq_credentials() + + self.assertIsInstance(rabbitmq_creds, dict) + + def test_rabbitmq_envs_values(self): + rabbitmq_creds = load_postgres_credentials() + + self.assertNotEqual(rabbitmq_creds["user"], None) + self.assertNotEqual(rabbitmq_creds["password"], None) + self.assertNotEqual(rabbitmq_creds["host"], None) + self.assertNotEqual(rabbitmq_creds["port"], None) + + self.assertIsInstance(rabbitmq_creds["user"], str) + self.assertIsInstance(rabbitmq_creds["password"], str) + self.assertIsInstance(rabbitmq_creds["host"], str) + self.assertIsInstance(rabbitmq_creds["port"], str) \ No newline at end of file diff --git a/tests/unit/test_load_retriever_hyperparameters.py b/tests/unit/test_load_retriever_hyperparameters.py index eadcbdc..1f9c2fa 100644 --- a/tests/unit/test_load_retriever_hyperparameters.py +++ b/tests/unit/test_load_retriever_hyperparameters.py @@ -1,7 +1,7 @@ import unittest from unittest.mock import patch -from retrievers.utils.load_hyperparams import load_hyperparams +from bot.retrievers.utils.load_hyperparams import load_hyperparams class TestLoadHyperparams(unittest.TestCase): diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/credentials.py b/utils/credentials.py new file mode 100644 index 0000000..02b3fb7 --- /dev/null +++ b/utils/credentials.py @@ -0,0 +1,55 @@ +import os + +from dotenv import load_dotenv + + +def load_postgres_credentials() -> dict[str, str]: + """ + load posgresql db credentials from .env + + Returns: + --------- + postgres_creds : dict[str, Any] + postgresql credentials + a dictionary representive of + `user`: str + `password` : str + `host` : str + `port` : int + """ + load_dotenv() + + postgres_creds = {} + + postgres_creds["user"] = os.getenv("POSTGRES_USER", "") + postgres_creds["password"] = os.getenv("POSTGRES_PASS", "") + postgres_creds["host"] = os.getenv("POSTGRES_HOST", "") + postgres_creds["port"] = os.getenv("POSTGRES_PORT", "") + + return postgres_creds + + +def load_rabbitmq_credentials() -> dict[str, str]: + """ + load rabbitmq credentials from .env + + Returns: + --------- + rabbitmq_creds : dict[str, Any] + rabbitmq credentials + a dictionary representive of + `user`: str + `password` : str + `host` : str + `port` : int + """ + load_dotenv() + + rabbitmq_creds = {} + + rabbitmq_creds["user"] = os.getenv("RABBIT_USER", "") + rabbitmq_creds["password"] = os.getenv("RABBIT_PASSWORD", "") + rabbitmq_creds["host"] = os.getenv("RABBIT_HOST", "") + rabbitmq_creds["port"] = os.getenv("RABBIT_PORT", "") + + return rabbitmq_creds \ No newline at end of file diff --git a/worker.py b/worker.py index fca2379..1104bbd 100644 --- a/worker.py +++ b/worker.py @@ -3,6 +3,8 @@ from tc_messageBroker.rabbit_mq.event import Event from tc_messageBroker.rabbit_mq.queue import Queue +from utils.credentials import load_rabbitmq_credentials + # TODO: Update according to our requirements def do_something(recieved_data): @@ -29,10 +31,10 @@ def job_recieve(broker_url, port, username, password): if __name__ == "__main__": - # TODO: read from .env - broker_url = "localhost" - port = 5672 - username = "root" - password = "pass" + rabbit_creds = load_rabbitmq_credentials() + username = rabbit_creds['user'] + password = rabbit_creds['password'] + broker_url = rabbit_creds['host'] + port = rabbit_creds['port'] job_recieve(broker_url, port, username, password)