linw1995

version: '3.5'
services:
  mongo:
    image: mongo
    ports:
      - "27017-27118:27017"
    command: mongod --replSet main

import subprocess as sp


def get_node_ports():
    ports = []
    idx = 1
    while True:
        try:
            output = sp.check_output(
                f"docker-compose port --index {idx} mongo 27017",
                shell=True,
                encoding="utf-8",
                stderr=sp.DEVNULL,
            )
        except sp.CalledProcessError:
            break

        port = output.split(":")[1]
        ports.append(int(port))
        idx += 1

    if not ports:
        raise RuntimeError("no node is running.")

    return ports

import json
import subprocess as sp
import tempfile

from utils import get_node_ports


def generate_init_js(node_ports):
    members = [
        {"_id": idx, "host": f"host.docker.internal:{port}"}
        for idx, port in enumerate(node_ports)
    ]
    initiate_parameters = {"_id": "main", "members": members}

    init_js = f"rs.initiate({json.dumps(initiate_parameters, indent=2)})"
    return init_js


def main():
    ports = get_node_ports()
    init_js = generate_init_js(ports)
    print("\x1b[32m$\x1b[0m cat init.js")
    print(init_js)

    container_name = sp.check_output(
        "docker-compose ps | grep mongo_1 | awk '{ print $1 }'",
        shell=True,
        encoding="utf-8",
    ).strip()
    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8") as f:
        f.write(init_js)
        f.flush()

        path = f.name
        sp.check_output(f"docker cp {path} {container_name}:/init.js", shell=True)
        print("\x1b[32m$\x1b[0m copy init.js")

    output = sp.check_output(
        f'docker exec {container_name} sh -c "mongo --port 27017 < /init.js"',
        shell=True,
        encoding="utf-8",
    )
    print("\x1b[32m$\x1b[0m initiate")
    print(output)


if __name__ == "__main__":
    main()

import asyncio
import logging
import socket

from unittest import mock

import pytest

from motor.motor_asyncio import AsyncIOMotorClient

from utils import get_node_ports


@pytest.fixture(scope="session")
def node_addresses():
    node_ports = get_node_ports()
    return [("localhost", port) for port in node_ports]


@pytest.fixture(autouse=True)
def hook_getaddrinfo():
    original_getaddrinfo = socket.getaddrinfo

    def container_2_host_address(host, port, family=0, type=0, proto=0, flags=0):
        """
        由于 MongoDB 会下发 replica 在容器中的地址，
        必须替换成在主机上的正确地址
        """
        if host == "host.docker.internal":
            logging.debug("container 2 host")
            host = "localhost"
        else:
            logging.debug("no change %r", (host, port))

        return original_getaddrinfo(host, port, family, type, proto, flags)

    with mock.patch("socket.getaddrinfo", new=container_2_host_address):
        yield


@pytest.fixture
def uri(node_addresses):
    nodes = ",".join(f"{host}:{port}" for host, port in node_addresses)
    return f"mongodb://{nodes}/?replicaSet=main"


@pytest.fixture
async def create_db(uri):
    async def creator():
        client = AsyncIOMotorClient(uri)
        db = client["default"]
        return db

    db = await creator()

    # 删除掉所有 collections
    col_names = await db.list_collection_names()
    for col_name in col_names:
        await db.drop_collection(col_name)

    return creator


async def get_latest_oplog_timestamp(client):
    """
    获取最新的 oplog 时间戳
    """
    cursor = client.local.oplog.rs.find().sort("$natural", -1).limit(-1)
    await cursor.fetch_next
    return cursor.next_object()["ts"]


@pytest.mark.asyncio
async def test_change_stream(create_db):
    db = await create_db()
    col_name = "boo"
    await db.create_collection(col_name)
    logging.debug("create collection")

    current_timestamp = await get_latest_oplog_timestamp(db.client)
    logging.debug("latest oplog timestamp %r", current_timestamp)

    async def wait_for_change():
        async with db[col_name].watch(
            pipeline=[{"$match": {"operationType": "insert"}}],
            start_at_operation_time=current_timestamp,
        ) as change_stream:
            # 使用 startAtOperationTime 参数，避免漏掉插入事件
            # https://docs.mongodb.com/manual/changeStreams/#start-time
            # 可要创建一个 bson.Timestamp 对象
            logging.debug("wait for change")
            async for change in change_stream:
                return change

    fut = asyncio.ensure_future(wait_for_change())

    data = {"username": "black"}
    await db[col_name].insert_one(data)
    logging.debug("insert new document %r", data)

    change = await fut
    logging.debug("received change %r", change)

    assert data == change["fullDocument"]

pytest test.py --log-cli-level debug