Skip to content

Commit ea9b077

Browse files
committed
Improved copy tests
1 parent 2f1354f commit ea9b077

1 file changed

Lines changed: 10 additions & 6 deletions

File tree

test/pg_test.rb

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -96,14 +96,18 @@ def test_type_map_binary
9696

9797
def test_copy_text
9898
embedding = Pgvector::Vector.new([1, 2, 3])
99+
half_embedding = Pgvector::HalfVector.new([1, 2, 3])
100+
binary_embedding = Pgvector::Bit.new([true, false, true])
99101
sparse_embedding = Pgvector::SparseVector.new([1, 2, 3])
100102
coder = PG::TextEncoder::CopyRow.new
101-
conn.copy_data("COPY pg_items (embedding, sparse_embedding) FROM STDIN", coder) do
102-
conn.put_copy_data([embedding, sparse_embedding])
103+
conn.copy_data("COPY pg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN", coder) do
104+
conn.put_copy_data([embedding, half_embedding, binary_embedding, sparse_embedding])
103105
end
104106
res = conn.exec("SELECT * FROM pg_items").first
105-
assert_equal embedding.to_a, res["embedding"].to_a
106-
assert_equal sparse_embedding.to_a, res["sparse_embedding"].to_a
107+
assert_equal [1, 2, 3], res["embedding"]
108+
assert_equal [1, 2, 3], res["half_embedding"].to_a
109+
assert_equal "101", res["binary_embedding"]
110+
assert_equal [1, 2, 3], res["sparse_embedding"].to_a
107111
end
108112

109113
def test_copy_binary
@@ -114,8 +118,8 @@ def test_copy_binary
114118
conn.put_copy_data([embedding.to_binary, sparse_embedding.to_binary])
115119
end
116120
res = conn.exec("SELECT * FROM pg_items").first
117-
assert_equal embedding.to_a, res["embedding"].to_a
118-
assert_equal sparse_embedding.to_a, res["sparse_embedding"].to_a
121+
assert_equal [1, 2, 3], res["embedding"]
122+
assert_equal [1, 2, 3], res["sparse_embedding"].to_a
119123
end
120124

121125
def conn

0 commit comments

Comments
 (0)