import pytest from dm.docs import search as docs_search def test_search_query_validation_rejects_ambiguous_or_unsafe_text(): for value in (None, "", " padded", "padded ", "bad\x00query", "x" * 2001): with pytest.raises(ValueError): docs_search._query_text(value) def test_search_query_validation_preserves_readable_multiline_queries(): query = "S001 合同\n交付条件" assert docs_search._query_text(query) == query def test_search_top_k_validation_rejects_bool_non_integer_and_out_of_range_values(): for value in (True, 1.5, "5", 0, -1, 101): with pytest.raises(ValueError): docs_search._top_k(value) assert docs_search._top_k(1) == 1 assert docs_search._top_k(100) == 100 def test_query_vector_requires_non_empty_finite_numeric_1d_data(): for value in ([], [[1.0, 2.0]], [1.0, float("nan")], [float("inf")], ["not-a-number"]): with pytest.raises(ValueError): docs_search._query_vector(value) vector = docs_search._query_vector([0, 1.5, -2]) assert vector.dtype.name == "float32" assert vector.tolist() == [0.0, 1.5, -2.0] def test_vector_score_requires_finite_numeric_values(): for value in (None, "bad", float("nan"), float("inf"), float("-inf")): with pytest.raises(ValueError): docs_search._vector_score(value) assert docs_search._vector_score("0.875") == 0.875 def test_search_row_validation_rejects_malformed_metadata(): valid = ("doc-1", "contract", "title", "S001", 0, "content", 0.8) assert docs_search._search_row(valid) == valid for value in ( None, ("too", "short"), ("doc", "type", "title", ["S001"], 0, "content", 0.8), ("doc", "type", "title", "S001", 0, None, 0.8), ("doc", "type", "title", "S001", 0, "content", float("nan")), ): with pytest.raises(ValueError): docs_search._search_row(value) def test_search_closes_cursor_and_connection_after_fetch(monkeypatch): class Cursor: def __init__(self): self.closed = False def execute(self, sql, params): self.params = params def fetchmany(self, size): assert size == 16 return [] def close(self): self.closed = True class Connection: def __init__(self): self.cur = Cursor() self.closed = False def cursor(self): return self.cur def close(self): self.closed = True connection = Connection() monkeypatch.setattr(docs_search, "embed_one", lambda query, is_query: [0.1, 0.2]) monkeypatch.setattr(docs_search, "connect_vec", lambda: connection) assert docs_search.search("S001 contract") == [] assert connection.cur.closed is True assert connection.closed is True def test_search_rejects_driver_results_beyond_the_candidate_budget(monkeypatch): class Cursor: def __init__(self): self.closed = False def execute(self, sql, params): assert "LIMIT %s" in sql assert params[-1] == 12 def fetchmany(self, size): assert size == 13 return [("unexpected",)] * size def close(self): self.closed = True class Connection: def __init__(self): self.cur = Cursor() self.closed = False def cursor(self): return self.cur def close(self): self.closed = True connection = Connection() monkeypatch.setattr(docs_search, "embed_one", lambda query, is_query: [0.1, 0.2]) monkeypatch.setattr(docs_search, "connect_vec", lambda: connection) with pytest.raises(RuntimeError, match="12-row candidate budget"): docs_search.search("S001 contract", top_k=1) assert connection.cur.closed is True assert connection.closed is True