|
|
@@ -13,7 +13,10 @@ |
|
|
|
# See the License for the specific language governing permissions and |
|
|
|
# limitations under the License. |
|
|
|
|
|
|
|
import atexit |
|
|
|
import hashlib |
|
|
|
import os |
|
|
|
import uuid |
|
|
|
from inspect import getcallargs |
|
|
|
|
|
|
|
from mock import Mock, patch |
|
|
@@ -27,23 +30,80 @@ from synapse.http.server import HttpServer |
|
|
|
from synapse.server import HomeServer |
|
|
|
from synapse.storage import PostgresEngine |
|
|
|
from synapse.storage.engines import create_engine |
|
|
|
from synapse.storage.prepare_database import prepare_database |
|
|
|
from synapse.storage.prepare_database import ( |
|
|
|
_get_or_create_schema_state, |
|
|
|
_setup_new_database, |
|
|
|
prepare_database, |
|
|
|
) |
|
|
|
from synapse.util.logcontext import LoggingContext |
|
|
|
from synapse.util.ratelimitutils import FederationRateLimiter |
|
|
|
|
|
|
|
# set this to True to run the tests against postgres instead of sqlite. |
|
|
|
# It requires you to have a local postgres database called synapse_test, within |
|
|
|
# which ALL TABLES WILL BE DROPPED |
|
|
|
USE_POSTGRES_FOR_TESTS = False |
|
|
|
USE_POSTGRES_FOR_TESTS = os.environ.get("SYNAPSE_POSTGRES", False) |
|
|
|
POSTGRES_USER = os.environ.get("SYNAPSE_POSTGRES_USER", "postgres") |
|
|
|
POSTGRES_BASE_DB = "_synapse_unit_tests_base_%s" % (os.getpid(),) |
|
|
|
|
|
|
|
|
|
|
|
def setupdb(): |
|
|
|
|
|
|
|
# If we're using PostgreSQL, set up the db once |
|
|
|
if USE_POSTGRES_FOR_TESTS: |
|
|
|
pgconfig = { |
|
|
|
"name": "psycopg2", |
|
|
|
"args": { |
|
|
|
"database": POSTGRES_BASE_DB, |
|
|
|
"user": POSTGRES_USER, |
|
|
|
"cp_min": 1, |
|
|
|
"cp_max": 5, |
|
|
|
}, |
|
|
|
} |
|
|
|
config = Mock() |
|
|
|
config.password_providers = [] |
|
|
|
config.database_config = pgconfig |
|
|
|
db_engine = create_engine(pgconfig) |
|
|
|
db_conn = db_engine.module.connect(user=POSTGRES_USER) |
|
|
|
db_conn.autocommit = True |
|
|
|
cur = db_conn.cursor() |
|
|
|
cur.execute("DROP DATABASE IF EXISTS %s;" % (POSTGRES_BASE_DB,)) |
|
|
|
cur.execute("CREATE DATABASE %s;" % (POSTGRES_BASE_DB,)) |
|
|
|
cur.close() |
|
|
|
db_conn.close() |
|
|
|
|
|
|
|
# Set up in the db |
|
|
|
db_conn = db_engine.module.connect( |
|
|
|
database=POSTGRES_BASE_DB, user=POSTGRES_USER |
|
|
|
) |
|
|
|
cur = db_conn.cursor() |
|
|
|
_get_or_create_schema_state(cur, db_engine) |
|
|
|
_setup_new_database(cur, db_engine) |
|
|
|
db_conn.commit() |
|
|
|
cur.close() |
|
|
|
db_conn.close() |
|
|
|
|
|
|
|
def _cleanup(): |
|
|
|
db_conn = db_engine.module.connect(user=POSTGRES_USER) |
|
|
|
db_conn.autocommit = True |
|
|
|
cur = db_conn.cursor() |
|
|
|
cur.execute("DROP DATABASE IF EXISTS %s;" % (POSTGRES_BASE_DB,)) |
|
|
|
cur.close() |
|
|
|
db_conn.close() |
|
|
|
|
|
|
|
atexit.register(_cleanup) |
|
|
|
|
|
|
|
|
|
|
|
@defer.inlineCallbacks |
|
|
|
def setup_test_homeserver( |
|
|
|
name="test", datastore=None, config=None, reactor=None, **kargs |
|
|
|
cleanup_func, name="test", datastore=None, config=None, reactor=None, **kargs |
|
|
|
): |
|
|
|
"""Setup a homeserver suitable for running tests against. Keyword arguments |
|
|
|
are passed to the Homeserver constructor. If no datastore is supplied a |
|
|
|
datastore backed by an in-memory sqlite db will be given to the HS. |
|
|
|
""" |
|
|
|
Setup a homeserver suitable for running tests against. Keyword arguments |
|
|
|
are passed to the Homeserver constructor. |
|
|
|
|
|
|
|
If no datastore is supplied, one is created and given to the homeserver. |
|
|
|
|
|
|
|
Args: |
|
|
|
cleanup_func : The function used to register a cleanup routine for |
|
|
|
after the test. |
|
|
|
""" |
|
|
|
if reactor is None: |
|
|
|
from twisted.internet import reactor |
|
|
@@ -95,9 +155,11 @@ def setup_test_homeserver( |
|
|
|
kargs["clock"] = MockClock() |
|
|
|
|
|
|
|
if USE_POSTGRES_FOR_TESTS: |
|
|
|
test_db = "synapse_test_%s" % uuid.uuid4().hex |
|
|
|
|
|
|
|
config.database_config = { |
|
|
|
"name": "psycopg2", |
|
|
|
"args": {"database": "synapse_test", "cp_min": 1, "cp_max": 5}, |
|
|
|
"args": {"database": test_db, "cp_min": 1, "cp_max": 5}, |
|
|
|
} |
|
|
|
else: |
|
|
|
config.database_config = { |
|
|
@@ -107,6 +169,21 @@ def setup_test_homeserver( |
|
|
|
|
|
|
|
db_engine = create_engine(config.database_config) |
|
|
|
|
|
|
|
# Create the database before we actually try and connect to it, based off |
|
|
|
# the template database we generate in setupdb() |
|
|
|
if datastore is None and isinstance(db_engine, PostgresEngine): |
|
|
|
db_conn = db_engine.module.connect( |
|
|
|
database=POSTGRES_BASE_DB, user=POSTGRES_USER |
|
|
|
) |
|
|
|
db_conn.autocommit = True |
|
|
|
cur = db_conn.cursor() |
|
|
|
cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,)) |
|
|
|
cur.execute( |
|
|
|
"CREATE DATABASE %s WITH TEMPLATE %s;" % (test_db, POSTGRES_BASE_DB) |
|
|
|
) |
|
|
|
cur.close() |
|
|
|
db_conn.close() |
|
|
|
|
|
|
|
# we need to configure the connection pool to run the on_new_connection |
|
|
|
# function, so that we can test code that uses custom sqlite functions |
|
|
|
# (like rank). |
|
|
@@ -125,15 +202,35 @@ def setup_test_homeserver( |
|
|
|
reactor=reactor, |
|
|
|
**kargs |
|
|
|
) |
|
|
|
db_conn = hs.get_db_conn() |
|
|
|
# make sure that the database is empty |
|
|
|
if isinstance(db_engine, PostgresEngine): |
|
|
|
cur = db_conn.cursor() |
|
|
|
cur.execute("SELECT tablename FROM pg_tables where schemaname='public'") |
|
|
|
rows = cur.fetchall() |
|
|
|
for r in rows: |
|
|
|
cur.execute("DROP TABLE %s CASCADE" % r[0]) |
|
|
|
yield prepare_database(db_conn, db_engine, config) |
|
|
|
|
|
|
|
# Prepare the DB on SQLite -- PostgreSQL is a copy of an already up to |
|
|
|
# date db |
|
|
|
if not isinstance(db_engine, PostgresEngine): |
|
|
|
db_conn = hs.get_db_conn() |
|
|
|
yield prepare_database(db_conn, db_engine, config) |
|
|
|
db_conn.commit() |
|
|
|
db_conn.close() |
|
|
|
|
|
|
|
else: |
|
|
|
# We need to do cleanup on PostgreSQL |
|
|
|
def cleanup(): |
|
|
|
# Close all the db pools |
|
|
|
hs.get_db_pool().close() |
|
|
|
|
|
|
|
# Drop the test database |
|
|
|
db_conn = db_engine.module.connect( |
|
|
|
database=POSTGRES_BASE_DB, user=POSTGRES_USER |
|
|
|
) |
|
|
|
db_conn.autocommit = True |
|
|
|
cur = db_conn.cursor() |
|
|
|
cur.execute("DROP DATABASE IF EXISTS %s;" % (test_db,)) |
|
|
|
db_conn.commit() |
|
|
|
cur.close() |
|
|
|
db_conn.close() |
|
|
|
|
|
|
|
# Register the cleanup hook |
|
|
|
cleanup_func(cleanup) |
|
|
|
|
|
|
|
hs.setup() |
|
|
|
else: |
|
|
|
hs = HomeServer( |
|
|
|