在 python 中使用 docker 和 psycopg2 进行单元测试

Question

问题简要描述： Psycopg2 不会从 unittest 连接到 docker 内的测试数据库，但从控制台可以正常连接。

错误消息： psycopg2.OperationalError：服务器意外关闭了连接 这可能意味着服务器在处理请求之前或处理请求时异常终止。

详细信息： 我正在尝试在 docker 中设置一个测试数据库，该数据库将在测试之前创建并填充，然后在测试之后删除。

这是设置数据库的功能：

def set_up_empty_test_db():
    client = docker.from_env()
    try:
        testdb = client.containers.get("TestDB")
        testdb.stop()
        testdb.remove()
        testdb = client.containers.run(
            "postgres",
            ports={5432: 5433},
            detach=True,
            name="TestDB",
            environment=["POSTGRES_PASSWORD=yourPassword"],
        )
    except NotFound:
        testdb = client.containers.run(
            "postgres",
            ports={5432: 5433},
            detach=True,
            name="TestDB",
            environment=["POSTGRES_PASSWORD=yourPassword"],
        )

    while testdb.status != "running":
        testdb = client.containers.get("TestDB")
    return 

如果我从控制台启动此函数，它将正常工作，不会出现错误，并且我可以看到 TestDB 容器正在运行。我可以成功启动连接：

conn = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")

但在单元测试时却不起作用。这是测试代码：

class TestCreateCity(unittest.TestCase):
    def setUp(self):
        set_up_empty_test_db()
        con = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")

        self.assertIsNone(con.isolation_level)

        cur = con.cursor()
        sql_file = open(os.path.join(ROOT_DIR + "/ddl/creates/schema_y.sql"), "r")
        cur.execute(sql_file.readline())
        con.commit()
        con.close()
        self.session = Session(creator=con)

    def test_create_city(self):
        cs.create_city("Test_CITY", "US")
        q = self.session.query(City).filter(City.city_name == "Test_CITY").one()
        self.assertIs(q)
        self.assertEqual(q.city_name, "Test_CITY")
        self.assertEqual(q.country_code, "US")

尝试启动连接时它会中断。请指教。

Answer 1

我知道这是一个老问题，但今天我需要做同样的事情。启动 postgres 服务器后，您尝试过快地连接到它 - 这就是它在控制台中工作的原因。

您需要做的就是替换：

set_up_empty_test_db()
con = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")

和：

set_up_empty_test_db()
con = None
while con == None:
    try:
        con = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")
    except psycopg2.OperationalError:
        time.sleep(0.5);

希望这对其他人有帮助。干杯!