[feature]Finsh the influxdb connect

Potterluo · Potterluo · commit 8343813ef682 · 2025-10-30T19:23:46.000+08:00
diff --git a/test/.gitignore b/test/.gitignore
@@ -1,6 +1,8 @@
 reports/
 dataset/
 logs/
+.cache/
+backup/
 $null
 *__pycache__/
 .*
diff --git a/test/README_zh.md b/test/README_zh.md
@@ -6,7 +6,7 @@
 
 - [x] 🏗️ **多级别测试**: UnitTest(0) → Smoke(1) → Feature(2) → E2E(3)
 - [x] 🏷️ **灵活标记**: 支持功能标签、平台标签和可靠性标签
-- [x] 📊 **数据收集**: 集成InfluxDB性能数据推送
+- [ ] 📊 **数据收集**: 集成InfluxDB性能数据推送
 - [x] 📋 **精美报告**: Allure测试报告集成，支持静态HTML和动态服务模式
 - [x] 🔧 **配置管理**: 基于YAML的灵活配置系统
 - [x] 🚀 **自动化**: 支持并行测试执行和自动清理
diff --git a/test/common/influxdb_utils.py b/test/common/influxdb_utils.py
@@ -1,58 +1,224 @@
-"""
-InfluxDB Data Push Utility
-Provides convenient InfluxDB data writing functionality
-"""
-
-from datetime import datetime
-from typing import Dict, Any, Optional, Union
+import json
+import time
+from datetime import datetime, timezone
+import threading
+from typing import Dict, Any, Optional
 from influxdb_client import InfluxDBClient, Point, WritePrecision
 from influxdb_client.client.write_api import SYNCHRONOUS
-from config_utils import config_utils as config_instance
+from common.config_utils import config_utils as config_instance
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+
+
 
 class InfluxDBUtils:
-    """InfluxDB Utility Class"""
+    """
+    InfluxDB Utility with Singleton pattern, health check, retry, and backup/resend logic.
+    It uses a synchronous write approach for simplicity in error handling.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        """Ensures only a single instance of InfluxDBUtils is created (Singleton)."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, url: str, token: str, org: str, bucket: str,
+                 backup_file: str = "failed_push.jsonl", retry_interval: int = 5, heartbeat_interval: int = 30):
+        # Prevent re-initialization
+        if hasattr(self, "_initialized") and self._initialized:
+            return
+
+        self.url = url
+        self.token = token
+        self.org = org
+        self.bucket = bucket
+        self.backup_file = backup_file
+        self.retry_interval = retry_interval
+        self.heartbeat_interval = heartbeat_interval
+        self._client: Optional[InfluxDBClient] = None
+        self._write_api = None
+        self._stop_heartbeat = threading.Event()
+
+        self._connect()
+        self._start_heartbeat()
+
+        self._initialized = True
+
+    def _connect(self):
+        """Try to establish connection to InfluxDB with retry mechanism."""
+        max_attempts = 3
+        for attempt in range(max_attempts):
+            try:
+                self._client = InfluxDBClient(url=self.url, token=self.token, org=self.org)
+                
+                # Check server health right after client initialization
+                self._client.ping() 
+                
+                self._write_api = self._client.write_api(write_options=SYNCHRONOUS)
+                # logger.info("Connected to InfluxDB successfully.")
+                
+                # Attempt to resend backed up metrics immediately after connection is established
+                self._resend_backed_up_metrics() 
+                return
+            except Exception as e:
+                logger.error(f"InfluxDB connection failed (attempt {attempt + 1}/{max_attempts}): {e}")
+                if attempt < max_attempts - 1:
+                    time.sleep(self.retry_interval)
+        
+        # If all attempts fail, raise an error
+        raise ConnectionError(f"Failed to connect to InfluxDB after {max_attempts} attempts.")
+
+
+    def _start_heartbeat(self):
+        """Start periodic health check."""
+        def heartbeat():
+            while not self._stop_heartbeat.is_set():
+                try:
+                    # Use ping for a light-weight connectivity check
+                    health = self._client.ping()
+                    if not health:
+                        logger.warning("InfluxDB health check failed.")
+                except Exception as e:
+                    # This could indicate a lost connection
+                    logger.warning(f"InfluxDB heartbeat error: {e}")
+                
+                # Wait for the next interval or until stopped
+                time.sleep(self.heartbeat_interval)
+                
+        threading.Thread(target=heartbeat, daemon=True).start()
+
+    def stop(self):
+        """Stop heartbeat and close the InfluxDB client connection."""
+        self._stop_heartbeat.set()
+        if self._client:
+            self._client.close()
+            logger.info("InfluxDB client connection closed.")
+
+    def _resend_backed_up_metrics(self):
+        """Attempts to resend metrics from the backup file and clears the file on success."""
+        if not os.path.exists(self.backup_file) or os.path.getsize(self.backup_file) == 0:
+            logger.debug("No backup file found or file is empty to resend.")
+            return
+
+        successful_resends = []
+        failed_resends = []
+
+        logger.info(f"Attempting to resend metrics from backup file: {self.backup_file}")
+
+        try:
+            with open(self.backup_file, "r", encoding="utf-8") as f:
+                lines = f.readlines()
 
-    def __init__(self):
-        """Initialize InfluxDB connection"""
-        self.config = config_instance.get_config("influxdb")
+            for line in lines:
+                try:
+                    data = json.loads(line)
+                    
+                    # Reconstruct Point object. Note: Using the original timestamp from backup.
+                    point = Point(data["measurement"])
+                    point.time(data["time"], WritePrecision.NS) 
+                    
+                    # Add tags, including a marker for resend
+                    point.tag("resend", "true") 
+                    for k, v in data.get("tags", {}).items():
+                        point.tag(k, v)
+                    
+                    # Add field
+                    point.field(data["field_key"], data["field_value"])
 
+                    # Attempt to write (synchronous)
+                    self._write_api.write(bucket=self.bucket, org=self.org, record=point)
+                    successful_resends.append(line)
+                except Exception as e:
+                    logger.error(f"Failed to resend backed-up metric: {e}. Data: {line.strip()}")
+                    failed_resends.append(line)
+                    # Stop resending on the first failure to avoid processing the whole file only to fail again
+                    break 
+            
+            # Rewrite the backup file only with failed lines
+            if successful_resends or failed_resends:
+                with open(self.backup_file, "w", encoding="utf-8") as f:
+                    for line in failed_resends:
+                        f.write(line)
+                
+                if not failed_resends:
+                    logger.info("Successfully resent all backed-up metrics and cleared the backup file.")
+                else:
+                    logger.warning(f"Resent {len(successful_resends)} metrics. {len(failed_resends)} metrics remain in backup.")
 
-# Global InfluxDB utility instance
-influxdb_utils = InfluxDBUtils()
+        except Exception as fe:
+            logger.error(f"Error reading or processing backup file: {fe}")
 
 
-def push_to_influx(measurement: str,
-                  value: Union[int, float, str],
-                  tags: Optional[Dict[str, str]] = None,
-                  fields: Optional[Dict[str, Union[int, float, str]]] = None,
-                  timestamp: Optional[datetime] = None) -> bool:
+    def push_metric(self, measurement: str, field_key: str, field_value: Any, tags: Optional[Dict[str, str]] = None):
+        """Public interface: pushes a single metric point to InfluxDB."""
+        # Using current UTC time in nanosecond precision
+        point = Point(measurement).time(datetime.now(timezone.utc), WritePrecision.NS)
 
-    return None
+        if tags:
+            for k, v in tags.items():
+                point.tag(k, v)
+        
+        point.field(field_key, field_value)
 
+        try:
+            self._write_api.write(bucket=self.bucket, org=self.org, record=point)
+            logger.debug(f"Pushed metric: {measurement}={field_value} with tags {tags}")
+        except Exception as e:
+            # If write fails, log error and trigger local backup
+            logger.error(f"Failed to push metric {measurement}: {e}. Backing up.")
+            self._backup_metric(measurement, field_key, field_value, tags)
 
-def push_test_metric(test_name: str,
-                     metric_name: str,
-                     value: Union[int, float],
-                     additional_tags: Optional[Dict[str, str]] = None) -> bool:
-    print("Push to InfluxDB, To be implemented.")
+    def _backup_metric(self, measurement, field_key, field_value, tags):
+        """Stores the failed metric locally in JSON Line format."""
+        data = {
+            # Storing the exact time of the push failure for later resending
+            "time": datetime.now(timezone.utc).isoformat(),
+            "measurement": measurement,
+            "field_key": field_key,
+            "field_value": field_value,
+            "tags": tags or {}
+        }
+        try:
+            # Ensure the directory exists
+            os.makedirs(os.path.dirname(self.backup_file), exist_ok=True)
+            with open(self.backup_file, "a", encoding="utf-8") as f:
+                f.write(json.dumps(data, ensure_ascii=False) + "\n")
+            logger.warning(f"Saved failed metric to backup file: {self.backup_file}")
+        except Exception as fe:
+            logger.error(f"Failed to write backup file: {fe}")
 
+# Instantiate the Singleton instance using configuration values
+influxdb_utils = InfluxDBUtils(
+    url=config_instance.get_nested_config("influxdb.url", "localhost"),
+    token = config_instance.get_nested_config("influxdb.token", os.getenv("INFLUXDB_TOKEN")),
+    org=config_instance.get_nested_config("influxdb.org", "ucm"),
+    bucket=config_instance.get_nested_config("influxdb.bucket", "ucm_test_bucket"),
+    backup_file=config_instance.get_nested_config("influxdb.offline_file", ".cache/influxdb_offline.jsonl")
+)
 
+# Example usage
 if __name__ == "__main__":
-    # Simple data push
-    push_to_influx("response_time", 0.123)
-
-    # Data push with tags
-    push_to_influx("accuracy", 0.95, {
-        "model": "v1.0",
-        "platform": "gpu",
-        "test_case": "classification"
-    })
-
-    # Test metric push
-    push_test_metric("test_calculation_accuracy", "calculation_time", 0.001, {
-        "feature": "accuracy"
-    })
-
-    # Data push with timestamp
-    from datetime import datetime
-    push_to_influx("memory_usage", 1024, {"test": "memory"}, timestamp=datetime.now())
+    logging.basicConfig(level=logging.INFO)
+    try:
+        # This will attempt to connect and resend any backed up metrics
+        influxdb_utils.push_metric("cpu_usage", "value", 0.88, {"host": "server1", "region": "us-west"})
+        influxdb_utils.push_metric("memory_free", "bytes", 1024 * 1024 * 512, {"host": "server2"})
+        
+        # Simulate a connection failure (if you change the URL to something invalid and run this)
+        # The metric will be backed up locally
+        # influxdb_utils.push_metric("fail_test", "value", 999, {"host": "test_server"})
+        
+    except ConnectionError as ce:
+        logger.error(f"Application cannot start: {ce}")
+    finally:
+        # Ensures clean shutdown
+        influxdb_utils.stop()
diff --git a/test/config.yaml b/test/config.yaml
@@ -7,7 +7,7 @@ reports:
     filename: "report.html"
     title: "UCM Pytest Test Report"
   allure:
-    enabled: true
+    enabled: false
     html_enable: true
     serve_mode: true  # 使用allure serve mode
     serve_host: "localhost"
@@ -22,12 +22,11 @@ log:
 
 # InfluxDB Configuration
 influxdb:
-  host: localhost
-  port: 8086
-  token: your-influxdb-token-here
-  org: your-organization
-  bucket: test-metrics
-  timeout: 10
+  url: ""
+  token: ""
+  org: "influxdb"
+  bucket: "ucm"
+  offline_file: ".cache/influxdb_offline.jsonl"
 
 # LLM Connection Configuration
 llm_connection:
diff --git a/test/suites/test_demo_function.py b/test/suites/test_demo_function.py
@@ -182,4 +182,11 @@ def test_mixed_steps_and_attachments():
         Status: PASSED
         """
         allure.attach(report_content, name="Test Report", 
-                     attachment_type=allure.attachment_type.TEXT)
+                     attachment_type=allure.attachment_type.TEXT)
+
+
+from common.influxdb_utils import influxdb_utils as influxdb_instance
+@pytest.mark.feature("influxdb")
+def test_influxdb_push():
+    """Example test case for pushing data to InfluxDB"""
+    influxdb_instance.push_metric("cpu_usage", "value", 3.0, {"host": "server1"})
