From d8288fabc5179472df7ab2eaebf5bec018a29f0e Mon Sep 17 00:00:00 2001
From: wd0517 <me@wangdi.ink>
Date: Wed, 10 Jul 2024 23:05:15 +0800
Subject: [PATCH] update tests

---
 tests/integrations/test_utils.py  | 14 +++++++-------
 tidb_vector/integrations/utils.py |  2 +-
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tests/integrations/test_utils.py b/tests/integrations/test_utils.py
index dbe1998..502148d 100644
--- a/tests/integrations/test_utils.py
+++ b/tests/integrations/test_utils.py
@@ -6,7 +6,7 @@
 
 def test_extract_info_from_column_definition():
     # Test case with dimension and distance metric
-    column_type = "VECTOR<FLOAT>(128)"
+    column_type = "VECTOR(128)"
     column_comment = "hnsw(distance=cosine)"
     expected_result = (128, "cosine")
     assert (
@@ -15,7 +15,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with dimension but no distance metric
-    column_type = "VECTOR<FLOAT>(256)"
+    column_type = "VECTOR(256)"
     column_comment = "some comment"
     expected_result = (256, None)
     assert (
@@ -24,7 +24,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with no dimension and no distance metric
-    column_type = "VECTOR<FLOAT>"
+    column_type = "VECTOR"
     column_comment = "another comment"
     expected_result = (None, None)
     assert (
@@ -33,7 +33,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with no dimension and no comment
-    column_type = "VECTOR<FLOAT>"
+    column_type = "VECTOR"
     column_comment = ""
     expected_result = (None, None)
     assert (
@@ -42,7 +42,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with dimension but no comment
-    column_type = "VECTOR<FLOAT>(256)"
+    column_type = "VECTOR(256)"
     column_comment = ""
     expected_result = (256, None)
     assert (
@@ -51,7 +51,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case without index type
-    column_type = "VECTOR<FLOAT>"
+    column_type = "VECTOR"
     column_comment = "distance=l2"
     expected_result = (None, "l2")
     assert (
@@ -60,7 +60,7 @@ def test_extract_info_from_column_definition():
     )
 
     # Test case with addition comment content
-    column_type = "VECTOR<FLOAT>(128)"
+    column_type = "VECTOR(128)"
     column_comment = "test, hnsw(distance=l2)"
     expected_result = (128, "l2")
     assert (
diff --git a/tidb_vector/integrations/utils.py b/tidb_vector/integrations/utils.py
index 2169dae..1784f72 100644
--- a/tidb_vector/integrations/utils.py
+++ b/tidb_vector/integrations/utils.py
@@ -90,7 +90,7 @@ def extract_info_from_column_definition(column_type, column_comment):
     """
     # Try to extract the dimension, which is optional.
     dimension_match = re.search(
-        r"VECTOR<FLOAT>(?:\((\d+)\))?", column_type, re.IGNORECASE
+        r"VECTOR(?:\((\d+)\))?", column_type, re.IGNORECASE
     )
     dimension = (
         int(dimension_match.group(1))