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