Python - модульный тест для сравнения вывода для двух функций

Question

Я пытаюсь выполнить модульный тест, чтобы проверить, совпадает ли количество двух запросов.

Ниже приведен мой код. Может ли кто-нибудь помочь с тем, как я могу выполнить модульный тест в Python, чтобы проверить, равны ли значения в обеих этих функциях.

base.py

import psycopg2

def q1():
  dwh_connection = psycopg2.connect(connection_details)
  cur = dwh_connection.cursor()
  cur.execute("select count(*) from table_1 limit 4")
  dwh_connection.close()

def q2():
  dwh_connection = psycopg2.connect(connection_details)
  cur = dwh_connection.cursor()
  cur.execute("select count(*) from table_2 limit 4")
  dwh_connection.close()

Спасибо.

Изменить:

Сообщение об ошибке:

File "/Users/PycharmProjects/unit/test_calc.py", line 10, in test_queries self.assertEqual(q1(),q2()) 

NameError: name 'q1' is not defined

Ниже приведен код моего юнит-теста

import unittest
import base import q1, q2


class TestStringMethods(unittest.TestCase):

    def test_queries(self):
        self.assertEqual(q1(),q2())

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

Код для base.py

import psycopg2

def q1():
   dwh_connection = psycopg2.connect(conn_details)
   cur = dwh_connection.cursor()
   query = "select count(*) from tble_1;"
   cur.execute(query)
   var = cur.fetchone()
   print (var[0])
   dwh_connection.close()

def q2():
   dwh_connection = psycopg2.connect(conn_details)
   cur = dwh_connection.cursor()
   query = "select count(*) from tble_2;"
   cur.execute(query)
   var = cur.fetchone()
   print (var[0])
   dwh_connection.close()

q1()
q2()

The above code works just fine, if executed separately.

Answer 1

Вы можете выполнить юнит-тест с чем-то вроде

import unittest

class TestExample(unittest.TestCase):

    def test_queries(self):
        self.assertEqual(q1(), q2())

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

Однако ваши функции сейчас ничего не возвращают. Измените их на что-то вроде

def q1():
   dwh_connection = psycopg2.connect(connection_details)
   cur = dwh_connection.cursor()
   result = cur.execute("""select count(*) from table_1 limit 4""").fetchone()
   dwh_connection.close()
   return result