Merge pull request #79 from harshraj22/registration_fix

Registration fix

src/auth/main.py

@@ -29,7 +29,8 @@ def register(auth_details: AuthDetails):
     cursor_object.execute(f"""INSERT INTO auth VALUES ('{auth_details.username}', '{hashed_password}', 'Free')""")
 
     database.commit()
-    requests.post(f'http://data_population:8020/sync/', json={'username': auth_details.username})
+    response = requests.post(f'http://data_population:8020/sync/', json={'username': auth_details.username})
+    print(f'Response from data_population for syncing data: {response.status_code}')
 
     logger.info(f"User {auth_details.username} registered successfully")
     return

src/data_population/populate.py

@@ -15,36 +15,51 @@ class Username(BaseModel):
 logger.setLevel(logging.DEBUG)
 
 
-database = mysql.connector.connect(
-  host ="subscription_db",
-  user ="root",
-  passwd ="mypassword",
-  database = 'DB'
-)
-
 app = FastAPI()
 
 r = redis.StrictRedis(host='redis', port=6379, decode_responses=True, password="")
 ALLOWED_PREFIX = 'allowed' # max number of requests allowed per WINDOW_LENGTH
 
-# preparing a cursor object
-cursor_object = database.cursor()
 
 @app.post('/sync', status_code=200)
 def sync(username: Username):
-  logger.info(f'Syncing {username} to redis')
+  database = mysql.connector.connect(
+    host ="subscription_db",
+    user ="root",
+    passwd ="mypassword",
+    database = 'DB'
+  )
+
+  # preparing a cursor object
+  cursor_object = database.cursor()
+
+  print(f'Syncing {username} to redis')
   user_record = f"""SELECT * FROM subscription_details where username = '{username.username}'"""
-  logger.info('User Record:', user_record)
+  print('User Record:', user_record)
   cursor_object.execute(user_record)
   result = cursor_object.fetchall()
 
   for username, subscription_tier, request_limit, retention_period in result:
     r.set(f'{ALLOWED_PREFIX}-{username}', request_limit) 
 
 
+  cursor_object.close()
+  database.close()
+
 @app.on_event("startup")
 @app.get('/sync_all', status_code=200)
 def sync_all():
+
+  database = mysql.connector.connect(
+    host ="subscription_db",
+    user ="root",
+    passwd ="mypassword",
+    database = 'DB'
+  )
+
+  # preparing a cursor object
+  cursor_object = database.cursor()
+
   user_record = """SELECT * FROM subscription_details """
   records = cursor_object.execute(user_record)
   result = cursor_object.fetchall()
@@ -57,6 +72,9 @@ def sync_all():
   for username, subscription_tier, request_limit, retention_period in result:
     r.set(f'{ALLOWED_PREFIX}-{username}', request_limit)
 
+  cursor_object.close()
+  database.close()
+
 @app.get('/healthcheck', status_code=200)
 def healthcheck():
   return 'OK'

src/frontend/locustfile.py

@@ -0,0 +1,53 @@
+import logging
+import numpy as np
+import base64
+import requests
+from locust import HttpUser, task, between
+
+API_URL = "http://0.0.0.0:8001/submit"
+LOGIN_URL = "http://0.0.0.0:8019/login"
+USERNAME = "test"
+PASSWORD = "test"
+
+logging.basicConfig(level=logging.INFO)
+
+def get_auth_token():
+    """
+    Retrieves authentication token by sending a POST request to login API.
+
+    Returns:
+        str: Authentication token.
+    """
+    data = {'username': USERNAME, 'password': PASSWORD}
+    response = requests.post(LOGIN_URL, json=data)
+    response.raise_for_status()
+    auth_token = response.json().get('token')
+    logging.info(f'Status_code: {response.status_code}, auth token: {auth_token}')
+    return auth_token
+
+img = np.random.rand(256, 256, 3)
+img = base64.b64encode(img).decode("utf-8")
+
+logging.info(f"curl -X POST -H 'Content-Type: application/json' -d '{{'username': '{USERNAME}', 'password': '{PASSWORD}'}}' {LOGIN_URL}")
+auth_token = get_auth_token()
+
+headers = {'User-Agent': 'Mozilla/5.0', 'Authorization': f'Bearer {auth_token}'}
+payload = {'img1': img, 'img2': img}
+
+def sync_health_check():
+    """
+    Sends a POST request to the API endpoint to perform a health check.
+    """
+    session = requests.Session()
+    response = session.post(API_URL, headers=headers, files=payload)
+    logging.info(f'Server responded with {response}')
+
+class QuickstartUser(HttpUser):
+    wait_time = between(1, 2)
+
+    @task
+    def check_api_endpoint(self):
+        """
+        Sends a POST request to the API endpoint using Locust client.
+        """
+        self.client.post("/submit", headers=headers, files=payload)

src/frontend/main_page.py

@@ -10,20 +10,21 @@
 
 if st.session_state.token is None:
     # Registration form
-    with st.form('Register'):
-        username = st.text_input('Username')
-        password = st.text_input('Password', type='password')
-        confirm_password = st.text_input('Confirm Password', type='password')
-        register = st.form_submit_button('Register')
-        if register:
-            if password == confirm_password:
-                r = requests.post('http://auth:8019/register', json={'username': username, 'password': password})
-                if r.status_code == 201:
-                    st.success('Registration successful')
+    with st.expander('Register New User'):
+        with st.form('Register'):
+            username = st.text_input('Username')
+            password = st.text_input('Password', type='password')
+            confirm_password = st.text_input('Confirm Password', type='password')
+            register = st.form_submit_button('Register')
+            if register:
+                if password == confirm_password:
+                    r = requests.post('http://auth:8019/register', json={'username': username, 'password': password})
+                    if r.status_code == 201:
+                        st.success('Registration successful')
+                    else:
+                        st.error('Registration failed')
                 else:
-                    st.error('Registration failed')
-            else:
-                st.error('Passwords do not match')
+                    st.error('Passwords do not match')
 
     with st.form('Authenticate'):
         username = st.text_input('Username')

src/frontend/requirements.txt

@@ -1 +1 @@
-streamlit==1.10.0
+streamlit==1.22.0

test/integration_test/main.py

@@ -8,8 +8,9 @@
 class TestImageSimilarityWorkflow(unittest.TestCase):
 
     def setUp(self):
+        r = requests.post('http://auth:8019/register', json={'username': 'free_user_007', 'password': 'free'})
         # Authenticate and store the auth token for later use
-        r = requests.post('http://auth:8019/login', json={'username': 'free_user', 'password': 'free'})
+        r = requests.post('http://auth:8019/login', json={'username': 'free_user_007', 'password': 'free'})
         self.assertEqual(r.status_code, 200)
         self.token = r.json()['token']