From f90e50343e492b4591407f546b9041f305e1f14a Mon Sep 17 00:00:00 2001 From: Hector Hernandez Date: Tue, 17 Mar 2020 16:05:15 -0700 Subject: [PATCH 1/8] WIP --- .../tests/docker-compose.yml | 11 +- .../tests/postgres/test_psycopg_functional.py | 103 ++++++++++++++++++ tox.ini | 3 + 3 files changed, 116 insertions(+), 1 deletion(-) create mode 100644 ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py diff --git a/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml b/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml index 1dab842f898..8f3ae1a3a1b 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml +++ b/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml @@ -4,4 +4,13 @@ services: otmongo: ports: - "27017:27017" - image: mongo:latest \ No newline at end of file + image: mongo:latest + + otpostgres: + image: postgres + ports: + - "5432:5432" + environment: + - POSTGRES_USER=postgres + - POSTGRES_PASSWORD=postgres + - POSTGRES_DB=postgres \ No newline at end of file diff --git a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py new file mode 100644 index 00000000000..bf54de8b69c --- /dev/null +++ b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py @@ -0,0 +1,103 @@ +# Copyright 2020, OpenTelemetry Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import typing +import unittest + +import psycopg2 + +from opentelemetry import trace as trace_api +from opentelemetry.ext.psycopg2 import trace_integration +from opentelemetry.sdk.trace import Span, Tracer, TracerProvider +from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +POSTGRES_HOST = os.getenv("POSTGRESQL_HOST ", "localhost") +POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT ", "5432")) +POSTGRES_DB_NAME = os.getenv("POSTGRESQL_DB_NAME ", "postgres") +POSTGRES_PASSWORD = os.getenv("POSTGRESQL_HOST ", "postgres") +POSTGRES_USER = os.getenv("POSTGRESQL_HOST ", "postgres") + + +class TestFunctionalPsycopg(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls._tracer_provider = TracerProvider() + cls._tracer = Tracer(cls._tracer_provider, None) + cls._span_exporter = InMemorySpanExporter() + cls._span_processor = SimpleExportSpanProcessor(cls._span_exporter) + cls._tracer_provider.add_span_processor(cls._span_processor) + trace_integration(cls._tracer) + cls._connection = psycopg2.connect( + dbname=POSTGRES_DB_NAME, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + host=POSTGRES_HOST, + port=POSTGRES_PORT, + ) + cls._cursor = cls._connection.cursor() + + def setUp(self): + self._span_exporter.clear() + + def validate_spans(self): + spans = self._span_exporter.get_finished_spans() + self.assertEqual(len(spans), 2) + for span in spans: + if span.name == "rootSpan": + root_span = span + else: + child_span = span + self.assertIsInstance(span.start_time, int) + self.assertIsInstance(span.end_time, int) + self.assertIsNot(root_span, None) + self.assertIsNot(child_span, None) + self.assertIsNotNone(child_span.parent) + self.assertEqual(child_span.parent.name, root_span.name) + self.assertIs(child_span.kind, trace_api.SpanKind.CLIENT) + # self.assertEqual( + # child_span.attributes["db.instance"], MONGODB_DB_NAME + # ) + # self.assertEqual( + # child_span.attributes["net.peer.name"], MONGODB_HOST + # ) + # self.assertEqual( + # child_span.attributes["net.peer.port"], MONGODB_PORT + # ) + + def test_execute(self): + """Should create a child span for execute method + """ + with self._tracer.start_as_current_span("rootSpan"): + self._cursor.execute("SELECT * FROM test") + self.validate_spans() + + # def test_executemany(self): + # """Should create a child span for update + # """ + # with self._tracer.start_as_current_span("rootSpan"): + # # self._collection.update_one( + # # {"name": "testName"}, {"$set": {"value": "someOtherValue"}} + # # ) + # self.validate_spans() + + # def test_callproc(self): + # """Should create a child span for find + # """ + # with self._tracer.start_as_current_span("rootSpan"): + # # self._collection.find_one() + # self.validate_spans() diff --git a/tox.ini b/tox.ini index 52c82eba075..edc4e77096c 100644 --- a/tox.ini +++ b/tox.ini @@ -241,6 +241,7 @@ deps = pytest docker-compose >= 1.25.2 pymongo ~= 3.1 + psycopg2 ~= 2.8.4 changedir = ext/opentelemetry-ext-docker-tests/tests @@ -248,6 +249,8 @@ changedir = commands_pre = pip install -e {toxinidir}/opentelemetry-api \ -e {toxinidir}/opentelemetry-sdk \ + -e {toxinidir}/ext/opentelemetry-ext-dbapi \ + -e {toxinidir}/ext/opentelemetry-ext-psycopg2 \ -e {toxinidir}/ext/opentelemetry-ext-pymongo - docker-compose up -d commands = From db8b1f528f6d14160cb7ef354796f8b37df41c0c Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Wed, 25 Mar 2020 13:07:40 -0700 Subject: [PATCH 2/8] Adding Postgres functional tests --- .../tests/docker-compose.yml | 11 +- .../tests/postgres/test_psycopg_functional.py | 128 ++++++++++++++++++ tox.ini | 3 + 3 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py diff --git a/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml b/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml index 1dab842f898..d2d844338df 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml +++ b/ext/opentelemetry-ext-docker-tests/tests/docker-compose.yml @@ -4,4 +4,13 @@ services: otmongo: ports: - "27017:27017" - image: mongo:latest \ No newline at end of file + image: mongo:latest + + otpostgres: + image: postgres + ports: + - "5432:5432" + environment: + POSTGRES_USER: testuser + POSTGRES_PASSWORD: testpassword + POSTGRES_DB: opentelemetry-tests \ No newline at end of file diff --git a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py new file mode 100644 index 00000000000..7cf3e000fec --- /dev/null +++ b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py @@ -0,0 +1,128 @@ +# Copyright 2020, OpenTelemetry Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import time +import unittest + +import psycopg2 + +from opentelemetry import trace as trace_api +from opentelemetry.ext.psycopg2 import trace_integration +from opentelemetry.sdk.trace import Tracer, TracerProvider +from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from pytest import fixture + +POSTGRES_HOST = os.getenv("POSTGRESQL_HOST ", "localhost") +POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT ", "5432")) +POSTGRES_DB_NAME = os.getenv("POSTGRESQL_DB_NAME ", "opentelemetry-tests") +POSTGRES_PASSWORD = os.getenv("POSTGRESQL_HOST ", "testpassword") +POSTGRES_USER = os.getenv("POSTGRESQL_HOST ", "testuser") + + +class TestFunctionalPsycopg(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls._connection = None + cls._cursor = None + cls._tracer_provider = TracerProvider() + cls._tracer = Tracer(cls._tracer_provider, None) + cls._span_exporter = InMemorySpanExporter() + cls._span_processor = SimpleExportSpanProcessor(cls._span_exporter) + cls._tracer_provider.add_span_processor(cls._span_processor) + trace_integration(cls._tracer) + + @classmethod + def tearDownClass(cls): + if cls._cursor: + cls._cursor.close() + if cls._connection: + cls._connection.close() + + @fixture(autouse=True) + def connect_to_db(self): + if self._connection is None: + # Try to connect to DB + for i in range(5): + try: + self._connection = psycopg2.connect( + dbname=POSTGRES_DB_NAME, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + host=POSTGRES_HOST, + port=POSTGRES_PORT, + ) + self._cursor = self._connection.cursor() + break + except Exception as err: + print(err) + time.sleep(5) + + def setUp(self): + self._span_exporter.clear() + + def validate_spans(self): + spans = self._span_exporter.get_finished_spans() + self.assertEqual(len(spans), 2) + for span in spans: + if span.name == "rootSpan": + root_span = span + else: + child_span = span + self.assertIsInstance(span.start_time, int) + self.assertIsInstance(span.end_time, int) + self.assertIsNot(root_span, None) + self.assertIsNot(child_span, None) + self.assertEqual(root_span.name, "rootSpan") + self.assertEqual(child_span.name, "postgresql.opentelemetry-tests") + self.assertIsNotNone(child_span.parent) + self.assertEqual(child_span.parent.name, root_span.name) + self.assertIs(child_span.kind, trace_api.SpanKind.CLIENT) + self.assertEqual( + child_span.attributes["db.instance"], POSTGRES_DB_NAME + ) + self.assertEqual(child_span.attributes["net.peer.name"], POSTGRES_HOST) + self.assertEqual(child_span.attributes["net.peer.port"], POSTGRES_PORT) + + def test_execute(self): + """Should create a child span for execute method + """ + with self._tracer.start_as_current_span("rootSpan"): + self._cursor.execute( + "CREATE TABLE IF NOT EXISTS test (id integer)" + ) + self._connection.commit() + self.validate_spans() + + def test_executemany(self): + """Should create a child span for executemany + """ + with self._tracer.start_as_current_span("rootSpan"): + data = ("1", "2", "3") + stmt = "INSERT INTO test (id) VALUES (%s)" + self._cursor.executemany(stmt, data) + self.validate_spans() + + def test_callproc(self): + """Should create a child span for callproc + """ + with self._tracer.start_as_current_span("rootSpan"): + try: + self._cursor.callproc("test", ()) + except Exception as err: + print(err) + self.validate_spans() diff --git a/tox.ini b/tox.ini index 52c82eba075..edc4e77096c 100644 --- a/tox.ini +++ b/tox.ini @@ -241,6 +241,7 @@ deps = pytest docker-compose >= 1.25.2 pymongo ~= 3.1 + psycopg2 ~= 2.8.4 changedir = ext/opentelemetry-ext-docker-tests/tests @@ -248,6 +249,8 @@ changedir = commands_pre = pip install -e {toxinidir}/opentelemetry-api \ -e {toxinidir}/opentelemetry-sdk \ + -e {toxinidir}/ext/opentelemetry-ext-dbapi \ + -e {toxinidir}/ext/opentelemetry-ext-psycopg2 \ -e {toxinidir}/ext/opentelemetry-ext-pymongo - docker-compose up -d commands = From 9893d3af72432d595af82b3cc8d7efe961fda0d0 Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Thu, 26 Mar 2020 10:25:20 -0700 Subject: [PATCH 3/8] Adding exception when db cannot connect --- .../tests/postgres/test_psycopg_functional.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py index 7cf3e000fec..aac75ad14a1 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py +++ b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py @@ -71,6 +71,8 @@ def connect_to_db(self): except Exception as err: print(err) time.sleep(5) + if self._connection is None: + raise Exception("Failed to connect to DB") def setUp(self): self._span_exporter.clear() From 173a450319cdd4fba304b6b7d25c852904652037 Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Fri, 27 Mar 2020 11:19:56 -0700 Subject: [PATCH 4/8] Using else --- .../tests/postgres/test_psycopg_functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py index aac75ad14a1..07a6997818b 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py +++ b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py @@ -71,7 +71,7 @@ def connect_to_db(self): except Exception as err: print(err) time.sleep(5) - if self._connection is None: + else: raise Exception("Failed to connect to DB") def setUp(self): From d56907d7a8b6bbc2c71df8e16461fd6fd34cf773 Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Fri, 27 Mar 2020 11:31:28 -0700 Subject: [PATCH 5/8] Addressing comments --- dev-requirements.txt | 3 --- tox.ini | 1 - 2 files changed, 4 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index 6ce54793067..3ffaba9d614 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -3,8 +3,5 @@ flake8~=3.7 isort~=4.3 black>=19.3b0,==19.* mypy==0.740 -sphinx~=2.1 -sphinx-rtd-theme~=0.4 -sphinx-autodoc-typehints~=1.10.2 pytest!=5.2.3 pytest-cov>=2.8 diff --git a/tox.ini b/tox.ini index edc4e77096c..33d213326de 100644 --- a/tox.ini +++ b/tox.ini @@ -197,7 +197,6 @@ commands = [testenv:docs] deps = - -c dev-requirements.txt -c docs-requirements.txt sphinx sphinx-rtd-theme From 0cb2bb4a6a1a97c78d7e691df242d3123cdd5dc2 Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Fri, 27 Mar 2020 11:32:53 -0700 Subject: [PATCH 6/8] Undo change made by mistake --- dev-requirements.txt | 3 +++ tox.ini | 1 + 2 files changed, 4 insertions(+) diff --git a/dev-requirements.txt b/dev-requirements.txt index 3ffaba9d614..6ce54793067 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -3,5 +3,8 @@ flake8~=3.7 isort~=4.3 black>=19.3b0,==19.* mypy==0.740 +sphinx~=2.1 +sphinx-rtd-theme~=0.4 +sphinx-autodoc-typehints~=1.10.2 pytest!=5.2.3 pytest-cov>=2.8 diff --git a/tox.ini b/tox.ini index 33d213326de..edc4e77096c 100644 --- a/tox.ini +++ b/tox.ini @@ -197,6 +197,7 @@ commands = [testenv:docs] deps = + -c dev-requirements.txt -c docs-requirements.txt sphinx sphinx-rtd-theme From 9556c18e56658a2c66c85002368dbf3b502ac7ed Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Fri, 3 Apr 2020 16:03:00 -0700 Subject: [PATCH 7/8] Addressing comments --- .../tests/check_availability.py | 56 +++++++++++++++++++ .../tests/postgres/test_psycopg_functional.py | 47 ++++++---------- tox.ini | 3 +- 3 files changed, 74 insertions(+), 32 deletions(-) create mode 100644 ext/opentelemetry-ext-docker-tests/tests/check_availability.py diff --git a/ext/opentelemetry-ext-docker-tests/tests/check_availability.py b/ext/opentelemetry-ext-docker-tests/tests/check_availability.py new file mode 100644 index 00000000000..e82a4dc50a6 --- /dev/null +++ b/ext/opentelemetry-ext-docker-tests/tests/check_availability.py @@ -0,0 +1,56 @@ +# Copyright 2020, OpenTelemetry Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import sys +import time + +import psycopg2 + +POSTGRES_HOST = os.getenv("POSTGRESQL_HOST ", "localhost") +POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT ", "5432")) +POSTGRES_DB_NAME = os.getenv("POSTGRESQL_DB_NAME ", "opentelemetry-tests") +POSTGRES_PASSWORD = os.getenv("POSTGRESQL_HOST ", "testpassword") +POSTGRES_USER = os.getenv("POSTGRESQL_HOST ", "testuser") + + +def check_postgres_connection(): + # Try to connect to DB + for i in range(5): + try: + connection = psycopg2.connect( + dbname=POSTGRES_DB_NAME, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + host=POSTGRES_HOST, + port=POSTGRES_PORT, + ) + connection.close() + break + except Exception as err: + print(err) + time.sleep(5) + else: + raise Exception("Failed to connect to Postgres DB") + + +def check_docker_services_availability(): + try: + print("Checking Postgres availability") + check_postgres_connection() + except Exception: + sys.exit(1) + sys.exit(0) + + +check_docker_services_availability() diff --git a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py index 07a6997818b..d58e332c1bc 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py +++ b/ext/opentelemetry-ext-docker-tests/tests/postgres/test_psycopg_functional.py @@ -25,7 +25,6 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, ) -from pytest import fixture POSTGRES_HOST = os.getenv("POSTGRESQL_HOST ", "localhost") POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT ", "5432")) @@ -45,6 +44,15 @@ def setUpClass(cls): cls._span_processor = SimpleExportSpanProcessor(cls._span_exporter) cls._tracer_provider.add_span_processor(cls._span_processor) trace_integration(cls._tracer) + cls._connection = psycopg2.connect( + dbname=POSTGRES_DB_NAME, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + host=POSTGRES_HOST, + port=POSTGRES_PORT, + ) + cls._connection.set_session(autocommit=True) + cls._cursor = cls._connection.cursor() @classmethod def tearDownClass(cls): @@ -53,27 +61,6 @@ def tearDownClass(cls): if cls._connection: cls._connection.close() - @fixture(autouse=True) - def connect_to_db(self): - if self._connection is None: - # Try to connect to DB - for i in range(5): - try: - self._connection = psycopg2.connect( - dbname=POSTGRES_DB_NAME, - user=POSTGRES_USER, - password=POSTGRES_PASSWORD, - host=POSTGRES_HOST, - port=POSTGRES_PORT, - ) - self._cursor = self._connection.cursor() - break - except Exception as err: - print(err) - time.sleep(5) - else: - raise Exception("Failed to connect to DB") - def setUp(self): self._span_exporter.clear() @@ -87,8 +74,8 @@ def validate_spans(self): child_span = span self.assertIsInstance(span.start_time, int) self.assertIsInstance(span.end_time, int) - self.assertIsNot(root_span, None) - self.assertIsNot(child_span, None) + self.assertIsNotNone(root_span) + self.assertIsNotNone(child_span) self.assertEqual(root_span.name, "rootSpan") self.assertEqual(child_span.name, "postgresql.opentelemetry-tests") self.assertIsNotNone(child_span.parent) @@ -107,7 +94,6 @@ def test_execute(self): self._cursor.execute( "CREATE TABLE IF NOT EXISTS test (id integer)" ) - self._connection.commit() self.validate_spans() def test_executemany(self): @@ -122,9 +108,8 @@ def test_executemany(self): def test_callproc(self): """Should create a child span for callproc """ - with self._tracer.start_as_current_span("rootSpan"): - try: - self._cursor.callproc("test", ()) - except Exception as err: - print(err) - self.validate_spans() + with self._tracer.start_as_current_span("rootSpan"), self.assertRaises( + Exception + ): + self._cursor.callproc("test", ()) + self.validate_spans() diff --git a/tox.ini b/tox.ini index edc4e77096c..250b9d7c4cf 100644 --- a/tox.ini +++ b/tox.ini @@ -252,7 +252,8 @@ commands_pre = -e {toxinidir}/ext/opentelemetry-ext-dbapi \ -e {toxinidir}/ext/opentelemetry-ext-psycopg2 \ -e {toxinidir}/ext/opentelemetry-ext-pymongo - - docker-compose up -d + docker-compose up -d + python check_availability.py commands = pytest {posargs} From 496f73960e332f0e46d03c72de8151bfc4e0fca6 Mon Sep 17 00:00:00 2001 From: Hector Hernandez Guzman Date: Mon, 6 Apr 2020 12:17:49 -0700 Subject: [PATCH 8/8] Addressing comments --- .../tests/check_availability.py | 59 +++++++++++++------ 1 file changed, 41 insertions(+), 18 deletions(-) diff --git a/ext/opentelemetry-ext-docker-tests/tests/check_availability.py b/ext/opentelemetry-ext-docker-tests/tests/check_availability.py index e82a4dc50a6..b0f1e55b208 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/check_availability.py +++ b/ext/opentelemetry-ext-docker-tests/tests/check_availability.py @@ -12,21 +12,47 @@ # See the License for the specific language governing permissions and # limitations under the License. import os -import sys import time +import traceback import psycopg2 +import pymongo -POSTGRES_HOST = os.getenv("POSTGRESQL_HOST ", "localhost") -POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT ", "5432")) -POSTGRES_DB_NAME = os.getenv("POSTGRESQL_DB_NAME ", "opentelemetry-tests") -POSTGRES_PASSWORD = os.getenv("POSTGRESQL_HOST ", "testpassword") -POSTGRES_USER = os.getenv("POSTGRESQL_HOST ", "testuser") +MONGODB_COLLECTION_NAME = "test" +MONGODB_DB_NAME = os.getenv("MONGODB_DB_NAME", "opentelemetry-tests") +MONGODB_HOST = os.getenv("MONGODB_HOST", "localhost") +MONGODB_PORT = int(os.getenv("MONGODB_PORT", "27017")) +POSTGRES_DB_NAME = os.getenv("POSTGRESQL_DB_NAME", "opentelemetry-tests") +POSTGRES_HOST = os.getenv("POSTGRESQL_HOST", "localhost") +POSTGRES_PASSWORD = os.getenv("POSTGRESQL_HOST", "testpassword") +POSTGRES_PORT = int(os.getenv("POSTGRESQL_PORT", "5432")) +POSTGRES_USER = os.getenv("POSTGRESQL_HOST", "testuser") +RETRY_COUNT = 5 +RETRY_INTERVAL = 5 # Seconds + + +def check_pymongo_connection(): + # Try to connect to DB + for i in range(RETRY_COUNT): + try: + client = pymongo.MongoClient( + MONGODB_HOST, MONGODB_PORT, serverSelectionTimeoutMS=2000 + ) + db = client[MONGODB_DB_NAME] + collection = db[MONGODB_COLLECTION_NAME] + collection.find_one() + client.close() + break + except Exception as ex: + if i == RETRY_COUNT - 1: + raise (ex) + traceback.print_exc() + time.sleep(RETRY_INTERVAL) def check_postgres_connection(): # Try to connect to DB - for i in range(5): + for i in range(RETRY_COUNT): try: connection = psycopg2.connect( dbname=POSTGRES_DB_NAME, @@ -37,20 +63,17 @@ def check_postgres_connection(): ) connection.close() break - except Exception as err: - print(err) - time.sleep(5) - else: - raise Exception("Failed to connect to Postgres DB") + except Exception as ex: + if i == RETRY_COUNT - 1: + raise (ex) + traceback.print_exc() + time.sleep(RETRY_INTERVAL) def check_docker_services_availability(): - try: - print("Checking Postgres availability") - check_postgres_connection() - except Exception: - sys.exit(1) - sys.exit(0) + # Check if Docker services accept connections + check_pymongo_connection() + check_postgres_connection() check_docker_services_availability()