IMDS Mock Server Testing

CMakeLists.txt

@@ -223,6 +223,9 @@ include (MakeDistFiles)
 
 # Enable CTest
 include (CTest)
+if (BUILD_TESTING)
+   include (TestFixtures)
+endif ()
 
 # Ensure the default behavior: don't ignore RPATH settings.
 set (CMAKE_SKIP_BUILD_RPATH OFF)

build/cmake/CMakeLists.txt

@@ -15,6 +15,7 @@ set (build_cmake_MODULES
    Sanitizers.cmake
    CCache.cmake
    LLDLinker.cmake
+   TestFixtures.cmake
 )
 
 set_local_dist (build_cmake_DIST_local

build/cmake/LoadTests.cmake

@@ -26,6 +26,12 @@ endif ()
 # Split lines on newlines
 string (REPLACE "\n" ";" lines "${tests_out}")
 
+# TODO: Allow individual test cases to specify the fixtures they want.
+set (all_fixtures "mongoc/fixtures/fake_imds")
+set (all_env
+    MCD_TEST_AZURE_IMDS_HOST=localhost:14987  # Refer: Fixtures.cmake
+    )
+
 # Generate the test definitions
 foreach (line IN LISTS lines)
     if (NOT line MATCHES "^/")
@@ -44,5 +50,7 @@ foreach (line IN LISTS lines)
         SKIP_REGULAR_EXPRESSION "@@ctest-skipped@@"
         # 45 seconds of timeout on each test.
         TIMEOUT 45
+        FIXTURES_REQUIRED "${all_fixtures}"
+        ENVIRONMENT "${all_env}"
         )
 endforeach ()

build/cmake/TestFixtures.cmake

@@ -0,0 +1,49 @@
+find_package (Python3 COMPONENTS Interpreter)
+
+if (NOT TARGET Python3::Interpreter)
+    message (STATUS "Python3 was not found, so test fixtures will not be defined")
+    return ()
+endif ()
+
+get_filename_component(_MONGOC_BUILD_SCRIPT_DIR "${CMAKE_CURRENT_LIST_DIR}" DIRECTORY)
+set (_MONGOC_PROC_CTL_COMMAND "$<TARGET_FILE:Python3::Interpreter>" -u -- "${_MONGOC_BUILD_SCRIPT_DIR}/proc-ctl.py")
+
+
+function (mongo_define_subprocess_fixture name)
+    cmake_parse_arguments(PARSE_ARGV 1 ARG "" "SPAWN_WAIT;STOP_WAIT;WORKING_DIRECTORY" "COMMAND")
+    string (MAKE_C_IDENTIFIER ident "${name}")
+    if (NOT ARG_SPAWN_WAIT)
+        set (ARG_SPAWN_WAIT 1)
+    endif ()
+    if (NOT ARG_STOP_WAIT)
+        set (ARG_STOP_WAIT 5)
+    endif ()
+    if (NOT ARG_WORKING_DIRECTORY)
+        set (ARG_WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}")
+    endif ()
+    if (NOT ARG_COMMAND)
+        message (SEND_ERROR "mongo_define_subprocess_fixture(${name}) requires a COMMAND")
+        return ()
+    endif ()
+    get_filename_component (ctl_dir "${CMAKE_CURRENT_BINARY_DIR}/${ident}.ctl" ABSOLUTE)
+    add_test (NAME "${name}/start"
+              COMMAND ${_MONGOC_PROC_CTL_COMMAND} start
+                "--ctl-dir=${ctl_dir}"
+                "--cwd=${ARG_WORKING_DIRECTORY}"
+                "--spawn-wait=${ARG_SPAWN_WAIT}"
+                -- ${ARG_COMMAND})
+    add_test (NAME "${name}/stop"
+              COMMAND ${_MONGOC_PROC_CTL_COMMAND} stop "--ctl-dir=${ctl_dir}" --if-not-running=ignore)
+    set_property (TEST "${name}/start" PROPERTY FIXTURES_SETUP "${name}")
+    set_property (TEST "${name}/stop" PROPERTY FIXTURES_CLEANUP "${name}")
+endfunction ()
+
+# Create a fixture that runs a fake Azure IMDS server
+mongo_define_subprocess_fixture(
+    mongoc/fixtures/fake_imds
+    SPAWN_WAIT 0.2
+    COMMAND
+        "$<TARGET_FILE:Python3::Interpreter>" -u --
+        "${_MONGOC_BUILD_SCRIPT_DIR}/bottle.py" fake_azure:imds
+            --bind localhost:14987  # Port 14987 chosen arbitrarily
+    )

build/fake_azure.py

@@ -1,19 +1,21 @@
-from __future__ import annotations
-
+import functools
+import json
 import sys
 import time
-import json
 import traceback
-import functools
+from pathlib import Path
+
 import bottle
-from bottle import Bottle, HTTPResponse, request
+from bottle import Bottle, HTTPResponse
 
 imds = Bottle(autojson=True)
 """An Azure IMDS server"""
 
-from typing import TYPE_CHECKING, Any, Callable, Iterable, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, cast, overload
 
-if TYPE_CHECKING:
+if not TYPE_CHECKING:
+    from bottle import request
+else:
     from typing import Protocol
 
     class _RequestParams(Protocol):
@@ -22,7 +24,7 @@ def __getitem__(self, key: str) -> str:
             ...
 
         @overload
-        def get(self, key: str) -> str | None:
+        def get(self, key: str) -> 'str | None':
             ...
 
         @overload
@@ -31,25 +33,30 @@ def get(self, key: str, default: str) -> str:
 
     class _HeadersDict(dict[str, str]):
 
-        def raw(self, key: str) -> bytes | None:
+        def raw(self, key: str) -> 'bytes | None':
             ...
 
     class _Request(Protocol):
-        query: _RequestParams
-        params: _RequestParams
-        headers: _HeadersDict
 
-    request: _Request
+        @property
+        def query(self) -> _RequestParams:
+            ...
 
+        @property
+        def params(self) -> _RequestParams:
+            ...
 
-def parse_qs(qs: str) -> dict[str, str]:
-    return dict(bottle._parse_qsl(qs))  # type: ignore
+        @property
+        def headers(self) -> _HeadersDict:
+            ...
 
+    request = cast('_Request', None)
 
-def require(cond: bool, message: str):
-    if not cond:
-        print(f'REQUIREMENT FAILED: {message}')
-        raise bottle.HTTPError(400, message)
+
+def parse_qs(qs: str) -> 'dict[str, str]':
+    # Re-use the bottle.py query string parser. It's a private function, but
+    # we're using a fixed version of Bottle.
+    return dict(bottle._parse_qsl(qs))  # type: ignore
 
 
 _HandlerFuncT = Callable[
@@ -58,6 +65,7 @@ def require(cond: bool, message: str):
 
 
 def handle_asserts(fn: _HandlerFuncT) -> _HandlerFuncT:
+    "Convert assertion failures into HTTP 400s"
 
     @functools.wraps(fn)
     def wrapped():
@@ -72,17 +80,10 @@ def wrapped():
     return wrapped
 
 
-def test_flags() -> dict[str, str]:
+def test_params() -> 'dict[str, str]':
     return parse_qs(request.headers.get('X-MongoDB-HTTP-TestParams', ''))
 
 
-def maybe_pause():
-    pause = int(test_flags().get('pause', '0'))
-    if pause:
-        print(f'Pausing for {pause} seconds')
-        time.sleep(pause)
-
-
 @imds.get('/metadata/identity/oauth2/token')
 @handle_asserts
 def get_oauth2_token():
@@ -91,10 +92,7 @@ def get_oauth2_token():
     resource = request.query['resource']
     assert resource == 'https://vault.azure.net', 'Only https://vault.azure.net is supported'
 
-    flags = test_flags()
-    maybe_pause()
-
-    case = flags.get('case')
+    case = test_params().get('case')
     print('Case is:', case)
     if case == '404':
         return HTTPResponse(status=404)
@@ -114,17 +112,18 @@ def get_oauth2_token():
     if case == 'slow':
         return _slow()
 
-    assert case is None or case == '', f'Unknown HTTP test case "{case}"'
+    assert case in (None, ''), 'Unknown HTTP test case "{}"'.format(case)
 
     return {
         'access_token': 'magic-cookie',
-        'expires_in': '60',
+        'expires_in': '70',
         'token_type': 'Bearer',
         'resource': 'https://vault.azure.net',
     }
 
 
 def _gen_giant() -> Iterable[bytes]:
+    "Generate a giant message"
     yield b'{ "item": ['
     for _ in range(1024 * 256):
         yield (b'null, null, null, null, null, null, null, null, null, null, '
@@ -136,6 +135,7 @@ def _gen_giant() -> Iterable[bytes]:
 
 
 def _slow() -> Iterable[bytes]:
+    "Generate a very slow message"
     yield b'{ "item": ['
     for _ in range(1000):
         yield b'null, '
@@ -144,6 +144,8 @@ def _slow() -> Iterable[bytes]:
 
 
 if __name__ == '__main__':
-    print(f'RECOMMENDED: Run this script using bottle.py in the same '
-          f'directory (e.g. [{sys.executable} bottle.py fake_azure:imds])')
+    print(
+        'RECOMMENDED: Run this script using bottle.py (e.g. [{} {}/bottle.py fake_azure:imds])'
+        .format(sys.executable,
+                Path(__file__).resolve().parent))
     imds.run()