TestQueue.py

import unittest
import subprocess
import os
import sys
import pymysql
from dotenv import load_dotenv, find_dotenv
import json


class TestQueue(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        # Set working directory to lambda-api
        os.chdir(os.path.dirname(os.path.abspath(__file__)) +
                 "/../../backend/lambda-api")
        # Load database credentials to os.environ

        load_dotenv('./.env')
        # Make database connection
        connection = pymysql.connect(host='pickmo.cm8t4da2mqrz.us-west-2.rds.amazonaws.com',
                                     user=os.environ['RDS_USERNAME'],
                                     password=os.environ['RDS_PWD'],
                                     db='pickmo_testdb')
        # Clear PickerQueue
        try:
            with connection.cursor() as cursor:
                sql = "DELETE FROM PickerQueue"
                cursor.execute(sql)
                connection.commit()
        # End connection
        finally:
            connection.close()

    def test_add_remove(self):
        username = 'new_user1'
        # Test adding picker to queue and then removing that picker
        load_dotenv('./.env')
        # Make database connection
        connection = pymysql.connect(host='pickmo.cm8t4da2mqrz.us-west-2.rds.amazonaws.com',
                                     user=os.environ['RDS_USERNAME'],
                                     password=os.environ['RDS_PWD'],
                                     db='pickmo_testdb')

        #-----------------------First Test-------------------------#
        # Invoke lambda code
        data = {
            'body': {
                'username': username
            },
            'headers': {
                'test': 'lambda'
            },
            'requestContext': {
                'authorizer': {
                    'claims': {
                        'email': username
                    }
                }
            }
        }
        bash_command = "sls invoke local -f queuePicker --data '" + \
            json.dumps(data) + "'"
        self.assertEqual(subprocess.check_call(bash_command, shell=True), 0)

        # Check database for newly added picker
        try:
            with connection.cursor() as cursor:
                sql = "SELECT * FROM PickerQueue WHERE username = '" + username + "'"
                cursor.execute(sql)    #
                result = cursor.fetchall()
                self.assertEqual(len(result), 1)
                self.assertEqual(result[0][0], username)
        finally:
            connection.close()

        #-----------------------Second Test-------------------------#
        data = {
            'pathParameters': {
                'username': username
            },
            'headers': {
                'test': 'lambda'
            },
            'requestContext': {
                'authorizer': {
                    'claims': {
                        'email': username
                    }
                }
            }
        }
        bash_command = "sls invoke local -f dequeuePicker --data '" + \
            json.dumps(data) + "'"
        self.assertEqual(subprocess.check_call(bash_command, shell=True), 0)

        connection = pymysql.connect(host='pickmo.cm8t4da2mqrz.us-west-2.rds.amazonaws.com',
                                     user=os.environ['RDS_USERNAME'],
                                     password=os.environ['RDS_PWD'],
                                     db='pickmo_testdb')

        # Check database for deleted picker
        try:
            with connection.cursor() as cursor:
                sql = "SELECT * FROM PickerQueue WHERE username = '" + username + "'"
                cursor.execute(sql)    #
                result = cursor.fetchall()
                self.assertEqual(len(result), 0)
        finally:
            connection.close()

    def test_error_duplicate_user(self):
        username = 'new_user1'
        # Test adding picker to queue and then removing that picker
        load_dotenv('./.env')
        # Make database connection
        connection = pymysql.connect(host='pickmo.cm8t4da2mqrz.us-west-2.rds.amazonaws.com',
                                     user=os.environ['RDS_USERNAME'],
                                     password=os.environ['RDS_PWD'],
                                     db='pickmo_testdb')

        # Invoke lambda code
        data = {
            'body': {
                'username': username
            },
            'headers': {
                'test': 'lambda'
            },
            'requestContext': {
                'authorizer': {
                    'claims': {
                        'email': username
                    }
                }
            }
        }
        # Add user
        bash_command = "sls invoke local -f queuePicker --data '" + \
            json.dumps(data) + "'"
        self.assertEqual(subprocess.check_call(bash_command, shell=True), 0)

        # Add user again
        bash_command = "sls invoke local -f queuePicker --data '" + \
            json.dumps(data) + "' > QueueRes.json"
        self.assertEqual(subprocess.check_call(bash_command, shell=True), 0)
        with open('QueueRes.json', 'r') as f:
            data = json.load(f)
        self.assertEqual(data["statusCode"], 400)
        # Remove the response file
        os.remove('QueueRes.json')

if __name__ == '__main__':
    unittest.main()