Day 3 of pytest

Author: billillib

days/10-12-pytest/code/db_object.py

@@ -0,0 +1,16 @@
+import psycopg2
+
+
+class MyDB:
+    def __init__(self):
+        self.conn = None
+
+    def __enter__(self):
+        self.conn = psycopg2.connect("host=localhost dbname=test user=postgres "
+                                     "password=docker")
+        self.conn.autocommit = True
+
+    def __exit__(self, *args):
+        if self.conn:
+            self.conn.close()
+            self.conn = None

days/10-12-pytest/code/hello.py

@@ -0,0 +1,31 @@
+from db_object import MyDB
+import psycopg2
+
+import pytest
+
+
+@pytest.fixture(scope='module')
+def cur():
+    print("creating connection")
+    db = MyDB()
+    db.conn = psycopg2.connect("host=localhost dbname=test user=postgres "
+                               "password=docker")
+    cur = db.conn.cursor()
+    yield cur
+    cur.close()
+    db.conn.close()
+    print("closing db")
+
+
+def test_table_count(cur):
+    cur.execute("SELECT COUNT(*) FROM test_table")
+    result = cur.fetchone()
+
+    assert int(result[0]) == 1
+
+
+def test_table_value(cur):
+    cur.execute("SELECT * FROM test_table")
+    result = cur.fetchone()
+
+    assert result == (1, 'a')