diff --git a/core/testcontainers/core/waiting_utils.py b/core/testcontainers/core/waiting_utils.py
index c7eaf72..f46d811 100644
--- a/core/testcontainers/core/waiting_utils.py
+++ b/core/testcontainers/core/waiting_utils.py
@@ -15,7 +15,7 @@
import re
import time
import traceback
-from typing import Any, Callable, Iterable, Mapping, Optional, TYPE_CHECKING, Union
+from typing import Any, Callable, Iterable, Mapping, Optional, TYPE_CHECKING, Union, TypeVar
import wrapt
from . import config
@@ -71,8 +71,11 @@ def wait_container_is_ready(*transient_exceptions) -> Callable:
return wrapper
+WT = TypeVar('WT')
+
+
@wait_container_is_ready()
-def wait_for(condition: Callable[..., bool]) -> bool:
+def wait_for(condition: Callable[..., WT]) -> WT:
return condition()
diff --git a/mongodb/tests/test_mongodb.py b/mongodb/tests/test_mongodb.py
index c778a01..0b33e93 100644
--- a/mongodb/tests/test_mongodb.py
+++ b/mongodb/tests/test_mongodb.py
@@ -1,3 +1,4 @@
+from pymongo.database import Database
from pymongo import MongoClient
from pymongo.errors import OperationFailure
import pytest
@@ -13,7 +14,7 @@ def test_docker_generic_db():
port = mongo_container.get_exposed_port(27017)
return MongoClient(f"mongodb://{host}:{port}")
- db = wait_for(connect).primer
+ db: Database = wait_for(connect).primer
result = db.restaurants.insert_one(
{
"address": {
