app/backend/app/flow/test_influxdb_node.py

"""
Test script for InfluxDbNode functionality.

This script demonstrates write and read operations with InfluxDB.
Requires an InfluxDB instance running.

To run a local InfluxDB instance:
    docker run -d -p 8086:8086 \
        -e DOCKER_INFLUXDB_INIT_MODE=setup \
        -e DOCKER_INFLUXDB_INIT_USERNAME=admin \
        -e DOCKER_INFLUXDB_INIT_PASSWORD=adminpass \
        -e DOCKER_INFLUXDB_INIT_ORG=my-org \
        -e DOCKER_INFLUXDB_INIT_BUCKET=sensors \
        -e DOCKER_INFLUXDB_INIT_ADMIN_TOKEN=my-super-secret-token \
        influxdb:2.7
"""

import asyncio
import sys
import time
from pathlib import Path

# Add flow directory to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from nodes import InfluxDbNode, Node
from pipeline import Pipeline
from state import MemoryState
from util import Message


# Default test configuration
TEST_CONFIG = {
    "url": "http://10.200.200.115:8086",
    "token": "***REMOVED-INFLUXDB-TOKEN***==",
    "org": "strobl",
    "bucket": "test",
}


def test_influxdb_node_validation():
    """Test InfluxDbNode parameter validation."""
    print("=" * 60)
    print("Testing InfluxDbNode Validation")
    print("=" * 60)

    # Test missing required params
    try:
        node = InfluxDbNode(
            requires=[Message(name="data", dtype=float)],
            params={"url": "http://localhost:8086"},  # Missing token, org, bucket
        )
        assert False, "Should have raised ValueError"
    except ValueError as e:
        print(f"✓ Correctly rejected missing params: {e}")

    # Test neither requires nor provides
    try:
        node = InfluxDbNode(params=TEST_CONFIG)
        assert False, "Should have raised ValueError"
    except ValueError as e:
        print(f"✓ Correctly rejected empty requires/provides: {e}")

    # Test valid write-only node
    writer = InfluxDbNode(
        requires=[Message(name="temperature", dtype=float)],
        params={
            **TEST_CONFIG,
            "writes": {
                "temperature": {
                    "measurement": "sensor_data",
                    "field": "temp_celsius",
                    "tags": {"location": "room1"},
                }
            },
        },
    )
    print(f"✓ Created write-only node: {writer.name}")
    assert len(writer.requires) == 1
    assert len(writer.provides) == 0

    # Test valid read-only node
    reader = InfluxDbNode(
        provides=[Message(name="avg_temperature", dtype=float)],
        params={
            **TEST_CONFIG,
            "queries": {
                "avg_temperature": {
                    "measurement": "sensor_data",
                    "field": "temp_celsius",
                    "aggregation": "mean",
                }
            },
        },
    )
    print(f"✓ Created read-only node: {reader.name}")
    assert len(reader.requires) == 0
    assert len(reader.provides) == 1

    # Test combined read/write node
    combined = InfluxDbNode(
        requires=[Message(name="raw_temp", dtype=float)],
        provides=[Message(name="avg_temp", dtype=float)],
        params={
            **TEST_CONFIG,
            "writes": {
                "raw_temp": {
                    "measurement": "temperature",
                    "field": "value",
                    "tags": {"source": "sensor"},
                }
            },
            "queries": {
                "avg_temp": {
                    "measurement": "temperature",
                    "field": "value",
                    "aggregation": "mean",
                }
            },
        },
    )
    print(f"✓ Created combined node: {combined.name}")
    assert len(combined.requires) == 1
    assert len(combined.provides) == 1

    print("\n✓ All validation tests passed!\n")


def test_flux_query_building():
    """Test Flux query string generation."""
    print("=" * 60)
    print("Testing Flux Query Building")
    print("=" * 60)

    node = InfluxDbNode(
        provides=[Message(name="value", dtype=float)],
        params={
            **TEST_CONFIG,
            "queries": {
                "value": {
                    "measurement": "temperature",
                    "field": "celsius",
                    "tags": {"location": "room1", "sensor": "dht22"},
                    "range": "-24h",
                    "aggregation": "mean",
                }
            },
        },
    )

    query = node._build_flux_query(
        measurement="temperature",
        field="celsius",
        tags={"location": "room1", "sensor": "dht22"},
        time_range="-24h",
        aggregation="mean",
    )

    print(f"Generated Flux query:\n{query}\n")

    # Verify query components
    assert 'from(bucket: "test")' in query
    assert "range(start: -24h)" in query
    assert 'r["_measurement"] == "temperature"' in query
    assert 'r["_field"] == "celsius"' in query
    assert 'r["location"] == "room1"' in query
    assert 'r["sensor"] == "dht22"' in query
    assert "mean()" in query

    print("✓ Query contains all expected components")

    # Test different aggregations
    for agg in ["last", "first", "max", "min", "sum", "count"]:
        query = node._build_flux_query(
            measurement="test",
            field="value",
            tags={},
            time_range="-1h",
            aggregation=agg,
        )
        assert f"{agg}()" in query
        print(f"✓ Aggregation '{agg}' works")

    print("\n✓ All query building tests passed!\n")


def test_influxdb_write(skip_if_no_server: bool = True):
    """Test writing data to InfluxDB."""
    print("=" * 60)
    print("Testing InfluxDB Write")
    print("=" * 60)

    # Create writer node with write configuration
    writer = InfluxDbNode(
        requires=[Message(name="temperature", dtype=float)],
        params={
            **TEST_CONFIG,
            "writes": {
                "temperature": {
                    "measurement": "test_temperature",
                    "field": "value",
                    "tags": {"location": "test_room", "sensor": "test_sensor"},
                }
            },
        },
        name="test_writer",
    )

    # Build a simple pipeline
    state = MemoryState()
    pipeline = Pipeline(nodes=[writer], state=state, max_workers=1)

    # Test data to write - now just a simple value!
    test_data = {"temperature": 25.5}

    print(f"Attempting to write: {test_data}")

    try:
        writer._write_points(test_data)
        print("✓ Write successful!")
    except Exception as e:
        if skip_if_no_server:
            print(f"⚠ Write failed (server may not be running): {e}")
        else:
            raise

    # Test with runtime tags override
    test_data_with_tags = {
        "temperature": {"value": 26.0, "tags": {"sensor": "override_sensor"}}
    }
    print(f"Attempting to write with runtime tags: {test_data_with_tags}")

    try:
        writer._write_points(test_data_with_tags)
        print("✓ Write with runtime tags successful!")
    except Exception as e:
        if skip_if_no_server:
            print(f"⚠ Write failed (server may not be running): {e}")
        else:
            raise

    print("\n✓ Write test completed!\n")


def test_influxdb_read(skip_if_no_server: bool = True):
    """Test reading data from InfluxDB."""
    print("=" * 60)
    print("Testing InfluxDB Read")
    print("=" * 60)

    # Create reader node
    reader = InfluxDbNode(
        provides=[
            Message(name="last_temp", dtype=float),
            Message(name="avg_temp", dtype=float),
        ],
        params={
            **TEST_CONFIG,
            "query_range": "-1h",
            "queries": {
                "last_temp": {
                    "measurement": "test_temperature",
                    "field": "value",
                    "tags": {"location": "test_room"},
                    "aggregation": "last",
                },
                "avg_temp": {
                    "measurement": "test_temperature",
                    "field": "value",
                    "tags": {"location": "test_room"},
                    "aggregation": "mean",
                },
            },
        },
        name="test_reader",
    )

    print(f"Attempting to query data...")

    try:
        results = reader._query_data()
        print(f"✓ Query successful! Results: {results}")

        if results:
            for key, value in results.items():
                print(f"  {key}: {value}")
        else:
            print("  No data found (this is OK if no data was written)")
    except Exception as e:
        if skip_if_no_server:
            print(f"⚠ Query failed (server may not be running): {e}")
        else:
            raise

    print("\n✓ Read test completed!\n")


def test_influxdb_pipeline_integration(skip_if_no_server: bool = True):
    """Test InfluxDbNode in a complete pipeline."""
    print("=" * 60)
    print("Testing InfluxDB Pipeline Integration")
    print("=" * 60)

    # Create a source node that generates sensor data (now just a float!)
    def generate_sensor_data(params, **kwargs):
        """Generate test sensor data."""
        import random

        return {"temperature": 20.0 + random.random() * 10}

    source = Node(
        f=generate_sensor_data,
        requires=[],
        provides=[Message(name="temperature", dtype=float)],
        params={},
        name="sensor_source",
    )

    # Create InfluxDB writer that receives the sensor data
    writer = InfluxDbNode(
        requires=[Message(name="temperature", dtype=float)],
        params={
            **TEST_CONFIG,
            "writes": {
                "temperature": {
                    "measurement": "pipeline_test",
                    "field": "temp_value",
                    "tags": {
                        "source": "test_pipeline",
                        "run_id": str(int(time.time())),
                    },
                }
            },
        },
        name="influx_writer",
    )

    # Build pipeline
    state = MemoryState()
    pipeline = Pipeline(
        nodes=[source, writer],
        state=state,
        max_workers=2,
    )

    print(f"Pipeline created with nodes: {[n.name for n in pipeline.nodes]}")
    print(f"Dependencies: {pipeline.dependencies}")

    try:
        # Trigger the source node to start the pipeline
        print("\nTriggering source node...")
        result = source.inject({})  # Source has no inputs, just generates output

        print(f"Pipeline execution result: {result}")
        print("✓ Pipeline integration successful!")
    except Exception as e:
        if skip_if_no_server:
            print(f"⚠ Pipeline test failed (server may not be running): {e}")
        else:
            raise

    print("\n✓ Pipeline integration test completed!\n")


async def test_influxdb_async_operations(skip_if_no_server: bool = True):
    """Test async write and query operations."""
    print("=" * 60)
    print("Testing InfluxDB Async Operations")
    print("=" * 60)

    # Create a combined node with write and query configs
    node = InfluxDbNode(
        requires=[Message(name="temperature", dtype=float)],
        provides=[Message(name="last_value", dtype=float)],
        params={
            **TEST_CONFIG,
            "writes": {
                "temperature": {
                    "measurement": "async_test",
                    "field": "value",
                    "tags": {"test": "async"},
                }
            },
            "queries": {
                "last_value": {
                    "measurement": "async_test",
                    "field": "value",
                    "aggregation": "last",
                }
            },
        },
        name="async_test_node",
    )

    try:
        # Test async write - now just a simple value!
        print("Testing async write...")
        await node.write_async({"temperature": 42.0})
        print("✓ Async write successful!")

        # Wait a moment for the write to be visible
        await asyncio.sleep(0.5)

        # Test async query
        print("Testing async query...")
        results = await node.query_async()
        print(f"✓ Async query successful! Results: {results}")

    except Exception as e:
        if skip_if_no_server:
            print(f"⚠ Async test failed (server may not be running): {e}")
        else:
            raise

    print("\n✓ Async operations test completed!\n")


def main():
    """Run all tests."""
    print("\n" + "=" * 60)
    print("InfluxDbNode Test Suite")
    print("=" * 60 + "\n")

    # Synchronous tests (always run)
    test_influxdb_node_validation()
    test_flux_query_building()

    # Tests that require InfluxDB server
    print("\n" + "-" * 60)
    print("The following tests require InfluxDB on localhost:8086")
    print("-" * 60 + "\n")

    try:
        test_influxdb_write(skip_if_no_server=True)
        test_influxdb_read(skip_if_no_server=True)
        test_influxdb_pipeline_integration(skip_if_no_server=True)
        asyncio.run(test_influxdb_async_operations(skip_if_no_server=True))
    except KeyboardInterrupt:
        print("\nTests interrupted by user")

    print("=" * 60)
    print("All tests completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()