from __future__ import annotations import io import json import os import subprocess import sys from pathlib import Path from typing import Any ORDERS_SOURCE = { "name": "orders", "table": "public.orders", "grain": ["id"], "columns": [ {"name": "id", "type": "number"}, {"name": "status", "type": "string"}, {"name": "amount", "type": "number"}, ], "joins": [], "measures": [{"name": "order_count", "expr": "count(*)"}], } def run_daemon_command( command: str, payload: dict[str, object] ) -> subprocess.CompletedProcess[str]: env = os.environ.copy() src_path = str(Path(__file__).resolve().parents[1] / "src") env["PYTHONPATH"] = src_path + os.pathsep + env.get("PYTHONPATH", "") return subprocess.run( [sys.executable, "-m", "ktx_daemon", command], input=json.dumps(payload), text=True, capture_output=True, check=False, env=env, ) def test_semantic_query_command_reads_stdin_and_writes_json() -> None: result = run_daemon_command( "semantic-query", { "sources": [ORDERS_SOURCE], "dialect": "postgres", "query": { "measures": ["orders.order_count"], "dimensions": ["orders.status"], }, }, ) assert result.returncode == 0, result.stderr parsed = json.loads(result.stdout) assert "public.orders" in parsed["sql"] assert parsed["columns"][0]["name"] == "orders.status" def test_semantic_validate_command_reads_stdin_and_writes_json() -> None: result = run_daemon_command( "semantic-validate", {"sources": [ORDERS_SOURCE], "dialect": "postgres"}, ) assert result.returncode == 0, result.stderr parsed = json.loads(result.stdout) assert parsed == { "valid": True, "errors": [], "warnings": [], "per_source_warnings": {}, } def test_command_returns_nonzero_for_invalid_json() -> None: env = os.environ.copy() src_path = str(Path(__file__).resolve().parents[1] / "src") env["PYTHONPATH"] = src_path + os.pathsep + env.get("PYTHONPATH", "") result = subprocess.run( [sys.executable, "-m", "ktx_daemon", "semantic-query"], input="{", text=True, capture_output=True, check=False, env=env, ) assert result.returncode == 1 assert "Expecting property name enclosed in double quotes" in result.stderr def test_serve_http_command_starts_uvicorn_without_reading_stdin( monkeypatch, ) -> None: from ktx_daemon import __main__ as daemon_main calls: list[dict[str, object]] = [] class FailingStdin: def read(self) -> str: raise AssertionError("serve-http must not read stdin JSON") def fake_run_http_server( *, host: str, port: int, log_level: str, enable_code_execution: bool, ) -> None: calls.append( { "host": host, "port": port, "log_level": log_level, "enable_code_execution": enable_code_execution, } ) monkeypatch.setattr(sys, "stdin", FailingStdin()) monkeypatch.setattr(daemon_main, "run_http_server", fake_run_http_server) assert ( daemon_main.main( [ "serve-http", "--host", "127.0.0.1", "--port", "9191", "--log-level", "warning", ] ) == 0 ) assert calls == [ { "host": "127.0.0.1", "port": 9191, "log_level": "warning", "enable_code_execution": False, } ] def test_serve_http_command_defaults_to_loopback(monkeypatch) -> None: from ktx_daemon import __main__ as daemon_main calls: list[dict[str, object]] = [] def fake_run_http_server( *, host: str, port: int, log_level: str, enable_code_execution: bool, ) -> None: calls.append( { "host": host, "port": port, "log_level": log_level, "enable_code_execution": enable_code_execution, } ) monkeypatch.setattr(daemon_main, "run_http_server", fake_run_http_server) assert daemon_main.main(["serve-http"]) == 0 assert calls == [ { "host": "127.0.0.1", "port": 8765, "log_level": "info", "enable_code_execution": False, } ] def test_serve_http_command_can_enable_code_execution(monkeypatch) -> None: from ktx_daemon import __main__ as daemon_main calls: list[dict[str, object]] = [] def fake_run_http_server( *, host: str, port: int, log_level: str, enable_code_execution: bool, ) -> None: calls.append( { "host": host, "port": port, "log_level": log_level, "enable_code_execution": enable_code_execution, } ) monkeypatch.setattr(daemon_main, "run_http_server", fake_run_http_server) assert daemon_main.main(["serve-http", "--enable-code-execution"]) == 0 assert calls == [ { "host": "127.0.0.1", "port": 8765, "log_level": "info", "enable_code_execution": True, } ] def test_lookml_parse_command_reads_stdin_and_writes_json() -> None: result = run_daemon_command( "lookml-parse", { "files": [ { "path": "views/orders.view.lkml", "content": """ view: orders { sql_table_name: public.orders ;; dimension: id { primary_key: yes type: number sql: ${TABLE}.id ;; } measure: order_count { type: count } } """, } ], "dialect": "postgres", }, ) assert result.returncode == 0, result.stderr parsed = json.loads(result.stdout) assert parsed["views"][0]["name"] == "orders" assert parsed["views"][0]["table_ref"] == "public.orders" assert parsed["views"][0]["measures"][0]["expr"] == "count(*)" assert parsed["joins"] == [] assert parsed["skipped_views"] == [] assert parsed["warnings"] == [] def test_semantic_generate_sources_command_reads_stdin_and_writes_json() -> None: result = run_daemon_command( "semantic-generate-sources", { "tables": [ { "name": "orders", "db": "public", "columns": [ {"name": "id", "type": "integer", "primary_key": True}, {"name": "amount", "type": "decimal"}, ], } ], "links": [], "dialect": "postgres", }, ) assert result.returncode == 0, result.stderr parsed = json.loads(result.stdout) assert parsed["source_count"] == 1 assert parsed["sources"][0]["name"] == "orders" assert parsed["sources"][0]["table"] == "public.orders" assert parsed["sources"][0]["measures"] == [ { "name": "record_count", "expr": "count(id)", "segments": [], "description": "Count of orders records", }, { "name": "total_amount", "expr": "sum(amount)", "segments": [], "description": "Sum of amount", }, { "name": "avg_amount", "expr": "avg(amount)", "segments": [], "description": "Average of amount", }, ] def test_database_introspect_command_reads_stdin_and_writes_json( monkeypatch, capsys ) -> None: from ktx_daemon import __main__ as daemon_main from ktx_daemon.database_introspection import ( DatabaseIntrospectionResponse, LiveDatabaseColumn, LiveDatabaseTable, ) def fake_introspect(request): assert request.connection_id == "warehouse" assert request.driver == "postgres" assert request.schemas == ["public"] assert request.table_scope is not None assert request.table_scope[0].db == "public" assert request.table_scope[0].name == "orders" return DatabaseIntrospectionResponse( connection_id="warehouse", extracted_at="2026-04-28T10:00:00+00:00", metadata={"driver": "postgres", "schemas": ["public"]}, tables=[ LiveDatabaseTable( catalog="warehouse", db="public", name="orders", columns=[ LiveDatabaseColumn( name="id", type="integer", nullable=False, primary_key=True, ) ], ) ], ) monkeypatch.setattr(daemon_main, "introspect_database_response", fake_introspect) monkeypatch.setattr( sys, "stdin", io.StringIO( '{"connection_id":"warehouse","driver":"postgres","url":"postgresql://readonly@example.test/warehouse","schemas":["public"],"table_scope":[{"db":"public","name":"orders"}]}' ), ) assert daemon_main.main(["database-introspect"]) == 0 captured = capsys.readouterr() parsed = json.loads(captured.out) assert parsed["connection_id"] == "warehouse" assert parsed["metadata"] == {"driver": "postgres", "schemas": ["public"]} assert parsed["tables"][0]["name"] == "orders" assert captured.err == "" def test_embedding_compute_command_reads_stdin_and_writes_json( monkeypatch, capsys ) -> None: from ktx_daemon import __main__ as daemon_main from ktx_daemon.embeddings import ComputeEmbeddingResponse def fake_compute(request): assert request.text == "hello" return ComputeEmbeddingResponse(embedding=[1.0, 2.0, 3.0]) monkeypatch.setattr(daemon_main, "compute_embedding_response", fake_compute) monkeypatch.setattr(sys, "stdin", io.StringIO('{"text": "hello"}')) assert daemon_main.main(["embedding-compute"]) == 0 captured = capsys.readouterr() assert json.loads(captured.out) == {"embedding": [1.0, 2.0, 3.0]} assert captured.err == "" def test_embedding_compute_bulk_command_reads_stdin_and_writes_json( monkeypatch, capsys ) -> None: from ktx_daemon import __main__ as daemon_main from ktx_daemon.embeddings import ComputeEmbeddingBulkResponse def fake_compute(request): assert request.texts == ["hello", "world"] return ComputeEmbeddingBulkResponse(embeddings=[[1.0, 2.0], [3.0, 4.0]]) monkeypatch.setattr(daemon_main, "compute_embedding_bulk_response", fake_compute) monkeypatch.setattr(sys, "stdin", io.StringIO('{"texts": ["hello", "world"]}')) assert daemon_main.main(["embedding-compute-bulk"]) == 0 captured = capsys.readouterr() assert json.loads(captured.out) == {"embeddings": [[1.0, 2.0], [3.0, 4.0]]} assert captured.err == "" def test_code_execute_command_reads_stdin_and_writes_json(monkeypatch, capsys) -> None: from ktx_daemon import __main__ as daemon_main from ktx_daemon.code_execution import ExecuteCodeResponse calls: list[dict[str, Any]] = [] def fake_execute(request, *, nest_api_url, auth_header): calls.append( { "request": request, "nest_api_url": nest_api_url, "auth_header": auth_header, } ) return ExecuteCodeResponse( formatted_result="\n\n=== Result ===\n\n7", result=7, ) monkeypatch.setattr(daemon_main, "execute_code_response", fake_execute) monkeypatch.setattr(sys, "stdin", io.StringIO('{"code": "result = 7"}')) assert daemon_main.main(["code-execute"]) == 0 captured = capsys.readouterr() assert json.loads(captured.out) == { "formatted_result": "\n\n=== Result ===\n\n7", "result": 7, "console_output": None, "error": None, "message": None, "visualizations": None, } assert captured.err == "" assert calls[0]["request"].code == "result = 7" assert calls[0]["nest_api_url"] is None assert calls[0]["auth_header"] is None