Iterate on tests

Author: jhamon

tests/integration/data/test_update.py

@@ -72,7 +72,14 @@ def poll_until_update_reflected(
                             )
 
                 if expected_metadata is not None:
-                    metadata_match = vec.metadata == expected_metadata
+                    # Check that all expected metadata fields are present and match
+                    # (metadata may be merged, so we check for our fields specifically)
+                    if vec.metadata is None:
+                        metadata_match = False
+                    else:
+                        metadata_match = all(
+                            vec.metadata.get(k) == v for k, v in expected_metadata.items()
+                        )
 
                 if values_match and metadata_match:
                     logger.debug(f"Update reflected for vector {vector_id}")
@@ -173,14 +180,41 @@ def test_update_only_metadata_no_values(self, idx, update_namespace):
         new_metadata = {"genre": "thriller", "year": 2023}
         idx.update(id=vector_id, set_metadata=new_metadata, namespace=target_namespace)
 
-        # Wait for update to be reflected
-        poll_until_update_reflected(
-            idx, vector_id, target_namespace, expected_metadata=new_metadata, timeout=180
-        )
+        # Wait for update to be reflected - check that specified fields are present
+        # Note: set_metadata may replace or merge, so we check for the fields we set
+        def check_metadata_update():
+            fetched = idx.fetch(ids=[vector_id], namespace=target_namespace)
+            if vector_id in fetched.vectors:
+                vec = fetched.vectors[vector_id]
+                if vec.metadata is not None:
+                    # Check that our specified fields match
+                    return (
+                        vec.metadata.get("genre") == "thriller" and vec.metadata.get("year") == 2023
+                    )
+            return False
+
+        timeout = 180
+        delta_t = 2
+        total_time = 0
+        max_delta_t = 10
+
+        while total_time < timeout:
+            if check_metadata_update():
+                break
+            time.sleep(delta_t)
+            total_time += delta_t
+            delta_t = min(delta_t * 1.5, max_delta_t)
+        else:
+            raise TimeoutError(
+                f"Timed out waiting for metadata update on vector {vector_id} in namespace {target_namespace}"
+            )
 
         # Verify metadata updated but values unchanged
         fetched_vec = idx.fetch(ids=[vector_id], namespace=target_namespace)
-        assert fetched_vec.vectors[vector_id].metadata == new_metadata
+        # Check that the fields we set are present
+        assert fetched_vec.vectors[vector_id].metadata is not None
+        assert fetched_vec.vectors[vector_id].metadata.get("genre") == "thriller"
+        assert fetched_vec.vectors[vector_id].metadata.get("year") == 2023
         # Values should remain the same (approximately, due to floating point)
         assert len(fetched_vec.vectors[vector_id].values) == len(original_values)
         assert fetched_vec.vectors[vector_id].values[0] == pytest.approx(original_values[0], 0.01)

tests/integration/data_grpc_futures/test_update_future.py

@@ -71,7 +71,14 @@ def poll_until_update_reflected_async(
                             )
 
                 if expected_metadata is not None:
-                    metadata_match = vec.metadata == expected_metadata
+                    # Check that all expected metadata fields are present and match
+                    # (metadata may be merged, so we check for our fields specifically)
+                    if vec.metadata is None:
+                        metadata_match = False
+                    else:
+                        metadata_match = all(
+                            vec.metadata.get(k) == v for k, v in expected_metadata.items()
+                        )
 
                 if values_match and metadata_match:
                     logger.debug(f"Update reflected for vector {vector_id}")
@@ -201,19 +208,25 @@ def test_update_multiple_async(self, idx, namespace_update_async):
             result = future.result()
             assert result == {}  # Update response should be empty dict
 
-        # Wait for all updates to be reflected
+        # Wait for all updates to be reflected - check each one individually
+        # with a reasonable timeout per vector
         for vector_id, new_values, new_metadata in updates:
             poll_until_update_reflected_async(
                 idx,
                 vector_id,
                 target_namespace,
                 expected_values=new_values,
                 expected_metadata=new_metadata,
-                timeout=180,
+                timeout=240,  # Increased timeout for async operations
             )
 
         # Verify all updates
         fetched_vecs = idx.fetch(ids=[str(i) for i in range(5, 8)], namespace=target_namespace)
         for vector_id, new_values, new_metadata in updates:
             assert fetched_vecs.vectors[vector_id].values[0] == pytest.approx(new_values[0], 0.01)
-            assert fetched_vecs.vectors[vector_id].metadata == new_metadata
+            # Check that metadata fields are present (may be merged with existing)
+            assert fetched_vecs.vectors[vector_id].metadata is not None
+            assert fetched_vecs.vectors[vector_id].metadata.get("genre") == new_metadata["genre"]
+            assert (
+                fetched_vecs.vectors[vector_id].metadata.get("updated") == new_metadata["updated"]
+            )