format; add a blob_type test to the vectorstore class

CassioML · May 23, 2024 · 86ad525 · 86ad525
1 parent e585975
commit 86ad525
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/tests/integration/test_tableclasses_metadatavectorcassandratable.py b/tests/integration/test_tableclasses_metadatavectorcassandratable.py
@@ -1,6 +1,7 @@
 """
 Table classes integration test - MetadataVectorCassandraTable
 """
+import datetime
 import math
 
 import pytest
@@ -199,3 +200,34 @@ async def test_crud_asyncio(self, db_session: Session, db_keyspace: str) -> None
         assert [r["row_id"] for r in ann_results_mt_l2] == ["theta_1", "theta_15"]
         assert ann_results_mt_l2[0]["distance"] < ann_results_mt_l2[1]["distance"]
         await t.aclear()
+
+    def test_body_type(self, db_session: Session, db_keyspace: str) -> None:
+        table_name = "m_v_ct_bt"
+        db_session.execute(f"DROP TABLE IF EXISTS {db_keyspace}.{table_name};")
+        #
+        t = MetadataVectorCassandraTable(
+            session=db_session,
+            keyspace=db_keyspace,
+            table=table_name,
+            vector_dimension=2,
+            primary_key_type="TEXT",
+            body_type="TIMESTAMP",
+        )
+
+        bblob = datetime.datetime(2024, 5, 23, 14, 58, 8)
+        t.put(
+            row_id="ts_id",
+            body_blob=bblob,
+            vector=[0.1, 0.2],
+            metadata={
+                "body_is_timestamp": True,
+            },
+        )
+
+        # reading
+        results = list(t.ann_search([0.11, 0.19], n=2))
+        assert results != []
+        result = results[0]
+        assert result["body_blob"] == bblob
+
+        db_session.execute(f"DROP TABLE IF EXISTS {db_keyspace}.{table_name};")
diff --git a/tests/integration/test_tableclasses_plaincassandratable.py b/tests/integration/test_tableclasses_plaincassandratable.py
@@ -69,8 +69,8 @@ def test_body_type(self, db_session: Session, db_keyspace: str) -> None:
             keyspace=db_keyspace,
             table=table_name,
             primary_key_type="TEXT",
-            body_type="INT"
+            body_type="INT",
         )
         t.put(row_id="row", body_blob=42)
         gotten = t.get(row_id="row")
-        assert gotten == {'row_id': 'row', 'body_blob': 42}
+        assert gotten == {"row_id": "row", "body_blob": 42}