refacto

Author: wartraxx51

python_testcontainers/infrahub_testcontainers/container.py

@@ -69,7 +69,9 @@ class InfrahubDockerCompose(DockerCompose):
     deployment_type: str | None = None
 
     @classmethod
-    def init(cls, directory: Path | None = None, version: str | None = None, deployment_type: str | None = None) -> Self:
+    def init(
+        cls, directory: Path | None = None, version: str | None = None, deployment_type: str | None = None
+    ) -> Self:
         if not directory:
             directory = Path.cwd()
 
@@ -112,7 +114,9 @@ def generate_project_name(cls) -> str:
 
     def create_docker_file(self, directory: Path) -> Path:
         current_directory = Path(__file__).resolve().parent
-        compose_file_name = "docker-compose-cluster.test.yml" if self.deployment_type == "cluster" else "docker-compose.test.yml"
+        compose_file_name = (
+            "docker-compose-cluster.test.yml" if self.deployment_type == "cluster" else "docker-compose.test.yml"
+        )
         compose_file = current_directory / compose_file_name
 
         test_compose_file = directory / "docker-compose.yml"
@@ -235,7 +239,7 @@ def database_create_backup(self, backup_name: str = "neo4j_database.backup", des
                 dest_dir / backup_name,
             )
 
-    def database_restore_backup(self, backup_file: Path) -> None:
+    def database_restore_backup(self, backup_file: Path) -> None:  # noqa: PLR0915
         assert self.use_neo4j_enterprise
 
         shutil.copy(
@@ -244,7 +248,7 @@ def database_restore_backup(self, backup_file: Path) -> None:
         )
         service_name = "database"
 
-        if self.deployment_type != "cluster":
+        if self.deployment_type != "cluster":  # noqa: PLR1702
             try:
                 self.get_container(service_name=service_name)
             except ContainerIsNotRunning:
@@ -351,8 +355,11 @@ def database_restore_backup(self, backup_file: Path) -> None:
 
             self.exec_in_container(
                 command=[
-                    "neo4j-admin", "database", "restore",
-                    "--from-path", str(self.internal_backup_dir / backup_file.name),
+                    "neo4j-admin",
+                    "database",
+                    "restore",
+                    "--from-path",
+                    str(self.internal_backup_dir / backup_file.name),
                     "neo4j",
                 ],
                 service_name=service_name,
@@ -361,42 +368,124 @@ def database_restore_backup(self, backup_file: Path) -> None:
             cmd = self.compose_command_property[:]
             cmd += ["restart", "database"]
             self._run_command(cmd=cmd)
-            # Waiting for cluster to stabilize...
-            time.sleep(10)
+
+            main_node = service_name
+            cluster_nodes = ["database", "database-core2", "database-core3"]
+
+            for attempt in range(3):
+                if attempt > 0:
+                    time.sleep(10)
+                    
+                (stdout, _, _) = self.exec_in_container(
+                    command=[
+                        "cypher-shell",
+                        "--format",
+                        "plain",
+                        "-d",
+                        "system",
+                        "-u",
+                        "neo4j",
+                        "-p",
+                        "admin",
+                        "SHOW DATABASES YIELD name, address, currentStatus WHERE name = 'system' RETURN address, currentStatus",
+                    ],
+                    service_name=main_node,
+                )
+
+                raw_output = stdout
+                nodes_status = dict.fromkeys(cluster_nodes, False)
+                online_count = 0
+                total_entries = 0
+                
+                try:
+                    for line_raw in stdout.splitlines():
+                        line = line_raw.strip()
+                        if not line or line.startswith("address"):
+                            continue
+                            
+                        total_entries += 1
+                        if "online" in line:
+                            online_count += 1
+                            for node in cluster_nodes:
+                                node_pattern = f'"{node}:'
+                                if node_pattern in line:
+                                    nodes_status[node] = True
+                                    break
+                    if all(nodes_status.values()) and online_count == len(cluster_nodes):
+                        break
+                except Exception as e:
+                    print(f"Error parsing database status on attempt {attempt + 1}: {e}")
+ 
+                print(f"Waiting for all nodes to be online. Current status: {nodes_status}")
+                time.sleep(5)
+            else:
+                debug_logs = {}
+                for node in cluster_nodes:
+                    try:
+                        (logs, _, _) = self.exec_in_container(
+                            command=["cat", "logs/debug.log"],
+                            service_name=node,
+                        )
+                        debug_logs[node] = logs
+                    except Exception as e:
+                        debug_logs[node] = f"Could not retrieve logs: {str(e)}"
+
+                debug_info = f"Raw output from SHOW DATABASES command:\n{raw_output}\n\n"
+                debug_info += f"Final node status: {nodes_status}\n\n"
+
+                status_str = ", ".join(
+                    [f"{node}: {'online' if status else 'offline'}" for node, status in nodes_status.items()]
+                )
+                logs_str = debug_info + "\n\n".join(
+                    [f"--- {node} logs ---\n{logs}" for node, logs in debug_logs.items()]
+                )
+
+                raise Exception(
+                    f"Failed to restore database cluster. Node status: {status_str}\nDebug logs:\n{logs_str}"
+                )
+                
 
             server_id = None
             try:
                 stdout, _, _ = self.exec_in_container(
                     command=[
-                        "cypher-shell", "-d", "system", "-u", "neo4j", "-p", "admin",
-                        'SHOW SERVERS YIELD name, address WHERE address = "database:7687" RETURN name;'
+                        "cypher-shell",
+                        "--format",
+                        "plain",
+                        "-d",
+                        "system",
+                        "-u",
+                        "neo4j",
+                        "-p",
+                        "admin",
+                        'SHOW SERVERS YIELD name, address WHERE address = "database:7687" RETURN name;',
                     ],
                     service_name=service_name,
                 )
-                print(f"Server ID query result: {stdout}")
 
-                lines = [line.strip() for line in stdout.splitlines() if line.strip() and not line.startswith("+")]
-                for line in lines:
-                    if line.startswith('"') and line.endswith('"'):
-                        server_id = line.strip('"')
-                        break
+                lines = stdout.splitlines()
+                for line_raw in lines:
+                    line = line_raw.strip()
+                    if not line or line == "name" or line.startswith("+"):
+                        continue
+                    server_id = line.strip('"')
+                    break
             except Exception as e:
                 print(f"Error retrieving server ID with direct query: {e}")
 
             if server_id:
-                print(f"Using server ID: {server_id} as seeder")
                 self.exec_in_container(
                     command=[
-                        "cypher-shell", "-d", "system", "-u", "neo4j", "-p", "admin",
-                        f"CREATE DATABASE neo4j TOPOLOGY 3 PRIMARIES OPTIONS {{ existingData: 'use', existingDataSeedInstance: '{server_id}' }};"
+                        "cypher-shell",
+                        "-d",
+                        "system",
+                        "-u",
+                        "neo4j",
+                        "-p",
+                        "admin",
+                        f"CREATE DATABASE neo4j TOPOLOGY 3 PRIMARIES OPTIONS {{ existingData: 'use', existingDataSeedInstance: '{server_id}' }};",
                     ],
                     service_name=service_name,
                 )
-
-            print("Checking database status...")
-            self.exec_in_container(
-                command=["cypher-shell", "-d", "system", "-u", "neo4j", "-p", "admin", "SHOW DATABASE neo4j;"],
-                service_name=service_name,
-            )
             self.start()
-            print("Database restored successfully")
+            print("Database restored successfully")

python_testcontainers/infrahub_testcontainers/helpers.py

@@ -62,7 +62,7 @@ def default_branch(self) -> str:
     @pytest.fixture(scope="class")
     def deployment_type(self, request: pytest.FixtureRequest) -> str | None:
         return request.config.getoption("infrahub_deployment_type")
-    
+
     @pytest.fixture(scope="class")
     def infrahub_compose(
         self,
@@ -72,7 +72,9 @@ def infrahub_compose(
         infrahub_version: str,
         deployment_type: str | None,
     ) -> InfrahubDockerCompose:
-        return InfrahubDockerCompose.init(directory=tmp_directory, version=infrahub_version, deployment_type=deployment_type)
+        return InfrahubDockerCompose.init(
+            directory=tmp_directory, version=infrahub_version, deployment_type=deployment_type
+        )
 
     @pytest.fixture(scope="class")
     def infrahub_app(self, request: pytest.FixtureRequest, infrahub_compose: InfrahubDockerCompose) -> dict[str, int]: