Replication Smoketest Cleanups

.github/workflows/ci.yml

@@ -60,8 +60,8 @@ jobs:
         with: { python-version: '3.12' }
         if: runner.os == 'Windows'
       - name: Run smoketests
-        # Note: clear_database only works in private
-        run: python -m smoketests ${{ matrix.smoketest_args }} -x clear_database
+        # Note: clear_database and replication only work in private
+        run: python -m smoketests ${{ matrix.smoketest_args }} -x clear_database replication
       - name: Stop containers (Linux)
         if: always() && runner.os == 'Linux'
         run: docker compose down

smoketests/__init__.py

@@ -6,7 +6,6 @@
 import re
 import shutil
 import string
-import string
 import subprocess
 import sys
 import tempfile
@@ -38,6 +37,9 @@
 # and a dotnet installation is detected
 HAVE_DOTNET = False
 
+# default value can be overriden by `--compose-file` flag
+COMPOSE_FILE = "./docker-compose.yml"
+
 # we need to late-bind the output stream to allow unittests to capture stdout/stderr.
 class CapturableHandler(logging.StreamHandler):
 
@@ -113,7 +115,7 @@ def run_cmd(*args, capture_stderr=True, check=True, full_output=False, cmd_name=
 
     needs_close = False
     if not capture_stderr:
-        logging.debug(f"--- stderr ---")
+        logging.debug("--- stderr ---")
         needs_close = True
 
     output = subprocess.run(
@@ -172,6 +174,12 @@ def call(self, reducer, *args, anon=False):
         anon = ["--anonymous"] if anon else []
         self.spacetime("call", *anon, "--", self.database_identity, reducer, *map(json.dumps, args))
 
+
+    def sql(self, sql):
+        self._check_published()
+        anon = ["--anonymous"]
+        return self.spacetime("sql", *anon, "--", self.database_identity, sql)
+
     def logs(self, n):
         return [log["message"] for log in self.log_records(n)]
 
@@ -181,6 +189,7 @@ def log_records(self, n):
         return list(map(json.loads, logs.splitlines()))
 
     def publish_module(self, domain=None, *, clear=True, capture_stderr=True):
+        print("publishing module", self.publish_module)
         publish_output = self.spacetime(
             "publish",
             *[domain] if domain is not None else [],
@@ -210,7 +219,7 @@ def subscribe(self, *queries, n):
         self._check_published()
         assert isinstance(n, int)
 
-        args = [SPACETIME_BIN, "--config-path", str(self.config_path),"subscribe", self.database_identity, "-t", "60", "-n", str(n), "--print-initial-update", "--", *queries]
+        args = [SPACETIME_BIN, "--config-path", str(self.config_path),"subscribe", self.database_identity, "-t", "600", "-n", str(n), "--print-initial-update", "--", *queries]
         fake_args = ["spacetime", *args[1:]]
         log_cmd(fake_args)
 
@@ -294,12 +303,12 @@ def tearDown(self):
 
     @classmethod
     def tearDownClass(cls):
-        if hasattr(cls, "database_identity"):
-            try:
-                # TODO: save the credentials in publish_module()
-                cls.spacetime("delete", cls.database_identity)
-            except Exception:
-                pass
+       if hasattr(cls, "database_identity"):
+           try:
+               # TODO: save the credentials in publish_module()
+               cls.spacetime("delete", cls.database_identity)
+           except Exception:
+               pass
 
     if sys.version_info < (3, 11):
         # polyfill; python 3.11 defines this classmethod on TestCase

smoketests/__main__.py

@@ -11,12 +11,13 @@
 import smoketests
 import sys
 import logging
+import itertools
 
 def check_docker():
     docker_ps = smoketests.run_cmd("docker", "ps", "--format=json")
     docker_ps = (json.loads(line) for line in docker_ps.splitlines())
     for docker_container in docker_ps:
-        if "node" in docker_container["Image"]:
+        if "node" in docker_container["Image"] or "spacetime" in docker_container["Image"]:
             return docker_container["Names"]
     else:
         print("Docker container not found, is SpacetimeDB running?")
@@ -51,13 +52,16 @@ def loadTestsFromName(self, name, module=None):
 def _convert_select_pattern(pattern):
     return f'*{pattern}*' if '*' not in pattern else pattern
 
+
 TESTPREFIX = "smoketests.tests."
 def main():
     tests = [fname.removesuffix(".py") for fname in os.listdir(TEST_DIR / "tests") if fname.endswith(".py") and fname != "__init__.py"]
 
     parser = argparse.ArgumentParser()
     parser.add_argument("test", nargs="*", default=tests)
     parser.add_argument("--docker", action="store_true")
+    parser.add_argument("--compose-file")
+    parser.add_argument("--no-docker-logs", action="store_true")
     parser.add_argument("--skip-dotnet", action="store_true", help="ignore tests which require dotnet")
     parser.add_argument("--show-all-output", action="store_true", help="show all stdout/stderr from the tests as they're running")
     parser.add_argument("--parallel", action="store_true", help="run test classes in parallel")
@@ -67,6 +71,7 @@ def main():
                         help='Only run tests which match the given substring')
     parser.add_argument("-x", dest="exclude", nargs="*", default=[])
     parser.add_argument("--no-build-cli", action="store_true", help="don't cargo build the cli")
+    parser.add_argument("--list", action="store_true", help="list the tests that would be run, but don't run them")
     args = parser.parse_args()
 
     if not args.no_build_cli:
@@ -94,9 +99,15 @@ def main():
     build_template_target()
 
     if args.docker:
-        docker_container = check_docker()
         # have docker logs print concurrently with the test output
-        subprocess.Popen(["docker", "logs", "-f", docker_container])
+        if args.compose_file:
+            smoketests.COMPOSE_FILE = args.compose_file
+        if not args.no_docker_logs:
+            if args.compose_file:
+                subprocess.Popen(["docker", "compose", "-f", args.compose_file, "logs", "-f"])
+            else:
+                docker_container = check_docker()
+                subprocess.Popen(["docker", "logs", "-f", docker_container])
         smoketests.HAVE_DOCKER = True
 
     smoketests.new_identity(TEST_DIR / 'config.toml')
@@ -116,6 +127,12 @@ def main():
     loader.testNamePatterns = args.testNamePatterns
 
     tests = loader.loadTestsFromNames(testlist)
+    if args.list:
+        print("Selected tests:\n")
+        for test in itertools.chain(*itertools.chain(*tests)):
+            print(f"{test}")
+        exit(0)
+
     buffer = not args.show_all_output
     verbosity = 2
 

smoketests/docker.py

@@ -0,0 +1,130 @@
+from dataclasses import dataclass
+import os
+import subprocess
+import time
+from typing import List, Optional
+from urllib.request import urlopen
+from . import COMPOSE_FILE
+import json
+
+def restart_docker():
+    docker = DockerManager(COMPOSE_FILE)
+    # Restart all containers.
+    docker.compose("restart")
+    # Ensure all nodes are reachable from outside.
+    containers = docker.list_containers()
+    for container in containers:
+        info = json.loads(docker._execute_command("docker", "inspect", container.name))
+        try:
+            port = info[0]['NetworkSettings']['Ports']['80/tcp'][0]['HostPort']
+        except KeyError:
+            continue
+        ping("127.0.0.1:{}".format(port))
+    # TODO: ping endpoint needs to wait for database startup & leader election
+    time.sleep(2)
+
+def ping(host):
+    tries = 0
+    while tries < 10:
+        tries += 1
+        try:
+            print(f"Ping Server at {host}")
+            urlopen(f"http://{host}/v1/ping")
+            print(f"Server up after {tries} tries")
+            break
+        except Exception:
+            print("Server down")
+            time.sleep(3)
+    else:
+        raise Exception(f"Server at {host} not responding")
+
+@dataclass
+class DockerContainer:
+    """Represents a Docker container with its basic properties."""
+    id: str
+    name: str
+
+class DockerManager:
+    """Manages all Docker and Docker Compose operations."""
+
+    def __init__(self, compose_file: str, **config):
+        self.compose_file = compose_file
+        self.network_name = config.get('network_name') or \
+                           os.getenv('DOCKER_NETWORK_NAME', 'private_spacetime_cloud')
+        self.control_db_container = config.get('control_db_container') or \
+                                  os.getenv('CONTROL_DB_CONTAINER', 'node')
+        self.spacetime_cli_bin = config.get('spacetime_cli_bin') or \
+                                os.getenv('SPACETIME_CLI_BIN', 'spacetimedb-cloud')
+
+    def _execute_command(self, *args: str) -> str:
+        """Execute a Docker command and return its output."""
+        try:
+            result = subprocess.run(
+                args,
+                capture_output=True,
+                text=True,
+                check=True
+            )
+            return result.stdout.strip()
+        except subprocess.CalledProcessError as e:
+            print(f"Command failed: {e.stderr}")
+            raise
+        except Exception as e:
+            print(f"Unexpected error: {str(e)}")
+            raise
+
+    def compose(self, *args: str) -> str:
+        """Execute a docker-compose command."""
+        return self._execute_command("docker", "compose", "-f", self.compose_file, *args)
+
+    def list_containers(self) -> List[DockerContainer]:
+        """List all containers and return as DockerContainer objects."""
+        output = self.compose("ps", "-a", "--format", "{{.ID}} {{.Name}}")
+        containers = []
+        for line in output.splitlines():
+            if line.strip():
+                container_id, name = line.split(maxsplit=1)
+                containers.append(DockerContainer(id=container_id, name=name))
+        return containers
+
+    def get_container_by_name(self, name: str) -> Optional[DockerContainer]:
+        """Find a container by name pattern."""
+        return next(
+            (c for c in self.list_containers() if name in c.name),
+            None
+        )
+
+    def kill_container(self, container_id: str):
+        """Kill a container by ID."""
+        print(f"Killing container {container_id}")
+        self._execute_command("docker", "kill", container_id)
+
+    def start_container(self, container_id: str):
+        """Start a container by ID."""
+        print(f"Starting container {container_id}")
+        self._execute_command("docker", "start", container_id)
+
+    def disconnect_container(self, container_id: str):
+        """Disconnect a container from the network."""
+        print(f"Disconnecting container {container_id}")
+        self._execute_command(
+            "docker", "network", "disconnect",
+            self.network_name, container_id
+        )
+        print(f"Disconnected container {container_id}")
+
+    def connect_container(self, container_id: str):
+        """Connect a container to the network."""
+        print(f"Connecting container {container_id}")
+        self._execute_command(
+            "docker", "network", "connect",
+            self.network_name, container_id
+        )
+        print(f"Connected container {container_id}")
+
+    def generate_root_token(self) -> str:
+        """Generate a root token using spacetimedb-cloud."""
+        return  self.compose(
+            "exec", self.control_db_container, self.spacetime_cli_bin, "token", "gen",
+            "--subject=placeholder-node-id",
+            "--jwt-priv-key", "/etc/spacetimedb/keys/id_ecdsa").split('|')[1]