feat: 完成 issue #4 数据总线 + 时序库 模板化封装
This commit is contained in:
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""测试引导：把 `core/data-bus` 以包名 `data_bus` 挂载到 sys.modules。
|
||||
|
||||
目录名 `data-bus` 含连字符，无法直接以包名 import；挂载后模块内相对导入
|
||||
（`from .templating import ...`）在 unittest 发现机制下可正常解析。
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
DATA_BUS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, DATA_BUS_DIR)
|
||||
if "data_bus" not in sys.modules:
|
||||
pkg = types.ModuleType("data_bus")
|
||||
pkg.__path__ = [DATA_BUS_DIR]
|
||||
sys.modules["data_bus"] = pkg
|
||||
@@ -0,0 +1,223 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""批量写入缓冲（batch_writer）测试：批量聚合 / 幂等去重 / 失败重试 / 端到端。
|
||||
|
||||
覆盖 Issue #4 验收点：端到端写入、批量写入（PRD 5.2 基线参数）、数据不丢不重。
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from data_bus.batch_writer import BatchWriter, MemorySink, TdengineSink
|
||||
from data_bus.templating import TemplateNaming
|
||||
|
||||
|
||||
class FakeClock:
|
||||
"""可控时钟：测试定时刷盘不依赖真实时间。"""
|
||||
|
||||
def __init__(self, start=0.0):
|
||||
self.now = start
|
||||
|
||||
def __call__(self):
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds):
|
||||
self.now += seconds
|
||||
|
||||
|
||||
class FlakySink(MemorySink):
|
||||
"""前 fail_times 次 write 抛异常，之后正常（模拟上游抖动）。"""
|
||||
|
||||
def __init__(self, fail_times=1):
|
||||
super().__init__()
|
||||
self.fail_times = fail_times
|
||||
self.failed_attempts = 0
|
||||
|
||||
def write(self, rows):
|
||||
if self.fail_times > 0:
|
||||
self.fail_times -= 1
|
||||
self.failed_attempts += 1
|
||||
raise RuntimeError("upstream unavailable")
|
||||
return super().write(rows)
|
||||
|
||||
|
||||
class PartialWriteSink(MemorySink):
|
||||
"""第一次 write 先落前 N 条再抛异常（模拟客户端部分写入后连接中断）。"""
|
||||
|
||||
def __init__(self, partial=3):
|
||||
super().__init__()
|
||||
self.partial = partial
|
||||
self.failed_attempts = 0
|
||||
|
||||
def write(self, rows):
|
||||
if self.failed_attempts == 0:
|
||||
self.failed_attempts += 1
|
||||
super().write(list(rows[: self.partial])) # 部分写入
|
||||
raise RuntimeError("connection lost after partial write")
|
||||
return super().write(rows)
|
||||
|
||||
|
||||
def make_row(i=0, point_id=None, ts=None):
|
||||
"""构造一条样本（对齐 edge-gateway 行格式 + batch 标签）。"""
|
||||
return {
|
||||
"device_id": f"CLF-{i % 5 + 1:02d}",
|
||||
"point_id": point_id or f"CLF-{i % 5 + 1:02d}.P{i:03d}",
|
||||
"value": 20.0 + i,
|
||||
"ts": float(ts) if ts is not None else 1754294400.0 + i,
|
||||
"quality": 1,
|
||||
"batch": f"b{int(i / 100)}",
|
||||
}
|
||||
|
||||
|
||||
class BatchFlushTest(unittest.TestCase):
|
||||
def test_flush_by_size(self):
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=250, flush_interval=999)
|
||||
for i in range(1000):
|
||||
writer.push(make_row(i))
|
||||
self.assertEqual(writer.stats()["flush_count"], 4)
|
||||
self.assertEqual(writer.stats()["written"], 1000)
|
||||
self.assertEqual(len(sink.rows), 1000)
|
||||
self.assertEqual(writer.pending(), 0)
|
||||
|
||||
def test_push_many_accepted_count(self):
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=100, flush_interval=999)
|
||||
accepted = writer.push_many([make_row(i) for i in range(50)])
|
||||
self.assertEqual(accepted, 50)
|
||||
self.assertEqual(writer.pending(), 50)
|
||||
|
||||
def test_close_flushes_pending(self):
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=1000, flush_interval=999)
|
||||
writer.push_many([make_row(i) for i in range(7)])
|
||||
writer.close()
|
||||
self.assertEqual(writer.stats()["written"], 7)
|
||||
self.assertTrue(sink.closed)
|
||||
|
||||
|
||||
class DedupTest(unittest.TestCase):
|
||||
def test_duplicate_push_rejected(self):
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=1000, flush_interval=999)
|
||||
row = make_row(0)
|
||||
self.assertTrue(writer.push(row))
|
||||
self.assertFalse(writer.push(dict(row))) # 同 设备-测点-ts → 幂等丢弃
|
||||
writer.flush()
|
||||
self.assertEqual(writer.stats()["duplicates"], 1)
|
||||
self.assertEqual(len(sink.rows), 1)
|
||||
|
||||
def test_dedup_across_batches(self):
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=50, flush_interval=999)
|
||||
rows = [make_row(i) for i in range(100)]
|
||||
writer.push_many(rows)
|
||||
# 跨批重推同键：仍只落一次
|
||||
re_pushed = writer.push_many(rows[:20])
|
||||
self.assertEqual(re_pushed, 0)
|
||||
writer.flush()
|
||||
self.assertEqual(writer.stats()["written"], 100)
|
||||
self.assertEqual(writer.stats()["duplicates"], 20)
|
||||
self.assertEqual(len(sink.rows), 100)
|
||||
|
||||
|
||||
class NoLossTest(unittest.TestCase):
|
||||
def test_push_survives_failed_auto_flush(self):
|
||||
# 自动刷盘失败不中断入队：失败批保留、下次 push 即自动重试（不丢）
|
||||
sink = FlakySink(fail_times=1)
|
||||
writer = BatchWriter(sink, batch_size=500, flush_interval=999)
|
||||
writer.push_many([make_row(i) for i in range(1000)])
|
||||
# 第 500 条触发失败刷盘（批保留），第 501 条 push 时重试成功 501 条
|
||||
self.assertEqual(writer.stats()["failed_flushes"], 1)
|
||||
self.assertEqual(writer.stats()["written"], 501)
|
||||
self.assertEqual(writer.pending(), 499) # 其余继续在缓冲，未丢失
|
||||
writer.flush() # 显式刷盘清空
|
||||
self.assertEqual(writer.stats()["written"], 1000)
|
||||
self.assertEqual(writer.pending(), 0)
|
||||
self.assertEqual(len(sink.rows), 1000) # 不丢
|
||||
self.assertEqual(sink.failed_attempts, 1)
|
||||
|
||||
def test_explicit_flush_retry(self):
|
||||
# 显式 flush 失败时缓冲保留，再次 flush 重试成功
|
||||
sink = FlakySink(fail_times=1)
|
||||
writer = BatchWriter(sink, batch_size=1000, flush_interval=999)
|
||||
writer.push_many([make_row(i) for i in range(50)])
|
||||
with self.assertRaises(RuntimeError):
|
||||
writer.flush()
|
||||
self.assertEqual(writer.pending(), 50) # 失败批保留
|
||||
writer.flush() # 重试成功
|
||||
self.assertEqual(writer.stats()["written"], 50)
|
||||
self.assertEqual(writer.pending(), 0)
|
||||
self.assertEqual(len(sink.rows), 50)
|
||||
|
||||
def test_no_dup_after_partial_retry(self):
|
||||
# sink 内部幂等：失败批次部分已写，重发也不会重复落库（不重）
|
||||
sink = PartialWriteSink(partial=3)
|
||||
writer = BatchWriter(sink, batch_size=10, flush_interval=999)
|
||||
writer.push_many([make_row(i) for i in range(20)])
|
||||
self.assertEqual(writer.stats()["failed_flushes"], 1)
|
||||
writer.flush()
|
||||
self.assertEqual(writer.pending(), 0)
|
||||
self.assertEqual(len(sink.rows), 20) # 无重复
|
||||
keys = [(r["device_id"], r["point_id"], r["ts"]) for r in sink.rows]
|
||||
self.assertEqual(len(keys), len(set(keys))) # 显式断言去重键唯一
|
||||
|
||||
|
||||
class TimeFlushTest(unittest.TestCase):
|
||||
def test_flush_after_interval(self):
|
||||
clock = FakeClock()
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=10000, flush_interval=0.1, clock=clock)
|
||||
writer.push_many([make_row(i) for i in range(5)])
|
||||
self.assertEqual(writer.pending(), 5) # 未达批大小且未超间隔
|
||||
clock.advance(0.1)
|
||||
writer.push(make_row(5)) # 超过 flush_interval → 自动刷盘
|
||||
self.assertEqual(writer.stats()["flush_count"], 1)
|
||||
self.assertEqual(writer.stats()["written"], 6)
|
||||
self.assertEqual(writer.pending(), 0)
|
||||
|
||||
|
||||
class EndToEndTest(unittest.TestCase):
|
||||
def test_memory_sink_round_trip(self):
|
||||
sink = MemorySink()
|
||||
writer = BatchWriter(sink, batch_size=5000, flush_interval=0.1)
|
||||
rows = [make_row(i) for i in range(300)]
|
||||
writer.push_many(rows)
|
||||
writer.flush()
|
||||
self.assertEqual(len(sink.rows), 300)
|
||||
# 保序：首条与末条一致
|
||||
self.assertEqual(sink.rows[0]["device_id"], rows[0]["device_id"])
|
||||
self.assertEqual(sink.rows[-1]["point_id"], rows[-1]["point_id"])
|
||||
# 标签字段透传（设备-测点-批次-质量）
|
||||
self.assertEqual({r["batch"] for r in sink.rows}, {"b0", "b1", "b2"})
|
||||
|
||||
def test_tdengine_sink_emits_batch_insert(self):
|
||||
executed = []
|
||||
naming = TemplateNaming(template="ti-cl4")
|
||||
sink = TdengineSink(naming, executor=lambda stmts: executed.extend(stmts))
|
||||
writer = BatchWriter(sink, batch_size=1000, flush_interval=999)
|
||||
# 同测点 20 条 → 单子表多值批量 INSERT
|
||||
rows = [
|
||||
{
|
||||
"device_id": "CLF-01",
|
||||
"point_id": "CLF-01.P000",
|
||||
"value": float(i),
|
||||
"ts": 1754294400.0 + i,
|
||||
"quality": 1,
|
||||
}
|
||||
for i in range(20)
|
||||
]
|
||||
writer.push_many(rows)
|
||||
writer.flush()
|
||||
self.assertEqual(sink.write_count, 20)
|
||||
self.assertEqual(len(executed), 1) # 一张子表 → 一条多值语句
|
||||
stmt = executed[0]
|
||||
self.assertTrue(stmt.startswith("INSERT INTO ti_cl4_pt_clf_01.p000 VALUES ("))
|
||||
self.assertEqual(stmt.count("), ("), 19) # 20 个值元组
|
||||
self.assertTrue(stmt.endswith(");"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,39 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""PostgreSQL schema 生成（postgres_schema）单元测试：DDL + 授权语句。"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from data_bus.postgres_schema import DEFAULT_TABLES, generate_grant_ddl, generate_schema_ddl
|
||||
from data_bus.templating import TemplateNaming
|
||||
|
||||
NAMING = TemplateNaming(template="ti-cl4")
|
||||
|
||||
|
||||
class SchemaDDLTest(unittest.TestCase):
|
||||
def test_schema_and_tables(self):
|
||||
ddl = generate_schema_ddl(NAMING)
|
||||
self.assertIn("CREATE SCHEMA IF NOT EXISTS tpl_ti_cl4;", ddl)
|
||||
for table in DEFAULT_TABLES:
|
||||
self.assertIn(f"CREATE TABLE IF NOT EXISTS tpl_ti_cl4.{table} (", ddl)
|
||||
|
||||
def test_models_foreign_key_uses_schema(self):
|
||||
ddl = generate_schema_ddl(NAMING)
|
||||
self.assertIn("REFERENCES tpl_ti_cl4.templates(template)", ddl)
|
||||
|
||||
|
||||
class GrantDDLTest(unittest.TestCase):
|
||||
def test_roles_ro_rw(self):
|
||||
grants = generate_grant_ddl(NAMING, roles=["databus"])
|
||||
text = "\n".join(grants)
|
||||
self.assertIn("GRANT USAGE ON SCHEMA tpl_ti_cl4 TO databus_ro;", text)
|
||||
self.assertIn("GRANT SELECT ON ALL TABLES IN SCHEMA tpl_ti_cl4 TO databus_ro;", text)
|
||||
self.assertIn("GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA tpl_ti_cl4 TO databus_rw;", text)
|
||||
self.assertIn("GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA tpl_ti_cl4 TO databus_rw;", text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,76 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""TDengine schema 生成（tdengine_schema）单元测试：DDL + 批量 INSERT。"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from data_bus.tdengine_schema import (
|
||||
PointSpec,
|
||||
build_batch_insert,
|
||||
generate_subtable_ddls,
|
||||
generate_supertable_ddl,
|
||||
)
|
||||
from data_bus.templating import TemplateNaming
|
||||
|
||||
NAMING = TemplateNaming(template="ti-cl4")
|
||||
|
||||
|
||||
def sample_point(device_id="CLF-01", point_id="CLF-01.TEMP", unit="℃", data_type="float"):
|
||||
return PointSpec(device_id=device_id, point_id=point_id, unit=unit, data_type=data_type)
|
||||
|
||||
|
||||
class SupertableDDLTest(unittest.TestCase):
|
||||
def test_ddl_shape(self):
|
||||
ddl = generate_supertable_ddl(NAMING)
|
||||
self.assertIn("CREATE STABLE IF NOT EXISTS ti_cl4_points", ddl)
|
||||
self.assertIn("ts TIMESTAMP", ddl)
|
||||
self.assertIn("device_id NCHAR(64)", ddl)
|
||||
|
||||
def test_retention_keep(self):
|
||||
ddl = generate_supertable_ddl(NAMING, retention_days=90)
|
||||
self.assertIn("KEEP(90)", ddl)
|
||||
|
||||
|
||||
class SubtableDDLTest(unittest.TestCase):
|
||||
def test_one_per_unique_point(self):
|
||||
points = [sample_point(), sample_point(), sample_point(point_id="CLF-01.PRES")]
|
||||
ddls = generate_subtable_ddls(NAMING, points)
|
||||
self.assertEqual(len(ddls), 2) # 同点位只生成一张子表
|
||||
self.assertTrue(any("ti_cl4_pt_clf_01.temp" in d for d in ddls))
|
||||
self.assertTrue(any("USING ti_cl4_points" in d for d in ddls))
|
||||
self.assertTrue(any("TAGS ('CLF-01', '℃', 'float')" in d for d in ddls))
|
||||
|
||||
|
||||
class BatchInsertTest(unittest.TestCase):
|
||||
def _row(self, point_id="CLF-01.TEMP", value=32.5, ts=1754294400.5, quality=1):
|
||||
return {
|
||||
"device_id": "CLF-01",
|
||||
"point_id": point_id,
|
||||
"value": value,
|
||||
"ts": ts,
|
||||
"quality": quality,
|
||||
}
|
||||
|
||||
def test_group_by_subtable_and_values(self):
|
||||
rows = [self._row(), self._row(value=33.1, ts=1754294401.0), self._row(point_id="CLF-01.PRES")]
|
||||
statements = build_batch_insert(NAMING, rows)
|
||||
self.assertEqual(len(statements), 2) # 两个测点 → 两张子表各一条多值 INSERT
|
||||
temp_stmt = [s for s in statements if "clf_01.temp" in s][0]
|
||||
self.assertIn(
|
||||
"INSERT INTO ti_cl4_pt_clf_01.temp VALUES (1754294400500, 32.5, 1), (1754294401000, 33.1, 1);",
|
||||
temp_stmt,
|
||||
)
|
||||
|
||||
def test_null_value(self):
|
||||
statements = build_batch_insert(NAMING, [self._row(value=None)])
|
||||
self.assertIn("NULL", statements[0])
|
||||
|
||||
def test_empty_rows(self):
|
||||
self.assertEqual(build_batch_insert(NAMING, []), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,89 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""模板化命名（templating）单元测试：Kafka / TDengine / PostgreSQL / MinIO。"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import _bootstrap # noqa: F401
|
||||
|
||||
from data_bus.templating import TemplateNaming, sanitize, sanitize_sql
|
||||
|
||||
|
||||
class SanitizeTest(unittest.TestCase):
|
||||
def test_lower_and_replace_unsafe(self):
|
||||
# 小写 + 空格/中文等非安全字符替换为 `_`，随后 strip 首尾 `._`
|
||||
self.assertEqual(sanitize("Ti-Cl4 模板"), "ti-cl4")
|
||||
self.assertEqual(sanitize("A.B-c_d"), "a.b-c_d")
|
||||
|
||||
def test_empty_fallback(self):
|
||||
self.assertEqual(sanitize(""), "tpl")
|
||||
self.assertEqual(sanitize("..."), "tpl")
|
||||
|
||||
def test_sql_extra_hyphen(self):
|
||||
# SQL 标识符额外把 `-` 换成 `_`，避免不加引号时报错
|
||||
self.assertEqual(sanitize_sql("ti-cl4"), "ti_cl4")
|
||||
self.assertEqual(sanitize_sql("A.B-c"), "a.b_c")
|
||||
|
||||
|
||||
class KafkaNamingTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.naming = TemplateNaming(template="ti-cl4", topic_prefix="ti-cl4", num_partitions=12)
|
||||
|
||||
def test_topic_format(self):
|
||||
self.assertEqual(self.naming.topic("CLF-01"), "ti-cl4.clf-01.points")
|
||||
|
||||
def test_partition_deterministic_and_in_range(self):
|
||||
p1 = self.naming.partition("CLF-01")
|
||||
p2 = self.naming.partition("CLF-01")
|
||||
self.assertEqual(p1, p2)
|
||||
self.assertGreaterEqual(p1, 0)
|
||||
self.assertLess(p1, 12)
|
||||
|
||||
def test_partitions_mapping(self):
|
||||
mapping = self.naming.partitions(["CLF-01", "CLF-02"])
|
||||
self.assertEqual(set(mapping), {"CLF-01", "CLF-02"})
|
||||
self.assertEqual(mapping["CLF-01"], self.naming.partition("CLF-01"))
|
||||
|
||||
|
||||
class TdengineNamingTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.naming = TemplateNaming(template="ti-cl4")
|
||||
|
||||
def test_stable_and_subtable(self):
|
||||
self.assertEqual(self.naming.stable(), "ti_cl4_points")
|
||||
# SQL 标识符保留 `.`（点位 ID 常用 `设备.测点` 形态）
|
||||
self.assertEqual(self.naming.subtable("CLF-01.TEMP"), "ti_cl4_pt_clf_01.temp")
|
||||
|
||||
|
||||
class PostgresNamingTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.naming = TemplateNaming(template="ti-cl4")
|
||||
|
||||
def test_schema_and_table(self):
|
||||
self.assertEqual(self.naming.pg_schema(), "tpl_ti_cl4")
|
||||
self.assertEqual(self.naming.pg_table("models"), "tpl_ti_cl4.models")
|
||||
|
||||
|
||||
class MinioNamingTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.naming = TemplateNaming(template="ti-cl4")
|
||||
|
||||
def test_bucket(self):
|
||||
self.assertEqual(self.naming.bucket(), "ti-cl4-artifacts")
|
||||
|
||||
def test_snapshot_key(self):
|
||||
self.assertEqual(
|
||||
self.naming.snapshot_key("quality-forecast", "2026-08-04", 3),
|
||||
"features/quality_forecast/2026-08-04/000003.jsonl",
|
||||
)
|
||||
|
||||
def test_model_artifact_key(self):
|
||||
self.assertEqual(
|
||||
self.naming.model_artifact_key("quality-forecast", "v1.2"),
|
||||
"models/quality_forecast/v1.2/model.bin",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user