Add Redis source code for testing

1 files changed, 36 insertions, 0 deletions

diff --git a/examples/redis-unstable/modules/vector-sets/tests/memory_usage.py b/examples/redis-unstable/modules/vector-sets/tests/memory_usage.py
new file mode 100644
index 0000000..d0f3f09
--- /dev/null
+++ b/examples/redis-unstable/modules/vector-sets/tests/memory_usage.py
@@ -0,0 +1,36 @@
+from test import TestCase, generate_random_vector
+import struct
+
+class MemoryUsageTest(TestCase):
+    def getname(self):
+        return "[regression] MEMORY USAGE with attributes"
+
+    def test(self):
+        # Generate random vectors
+        vec1 = generate_random_vector(4)
+        vec2 = generate_random_vector(4)
+        vec_bytes1 = struct.pack('4f', *vec1)
+        vec_bytes2 = struct.pack('4f', *vec2)
+
+        # Add vectors to the key, one with attribute, one without
+        self.redis.execute_command('VADD', self.test_key, 'FP32', vec_bytes1, f'{self.test_key}:item:1')
+        self.redis.execute_command('VADD', self.test_key, 'FP32', vec_bytes2, f'{self.test_key}:item:2', 'SETATTR', '{"color":"red"}')
+
+        # Get memory usage for the key
+        try:
+            memory_usage = self.redis.execute_command('MEMORY', 'USAGE', self.test_key)
+            # If we got here without exception, the command worked
+            assert memory_usage > 0, "MEMORY USAGE should return a positive value"
+
+            # Add more attributes to increase complexity
+            self.redis.execute_command('VSETATTR', self.test_key, f'{self.test_key}:item:1', '{"color":"blue","size":10}')
+
+            # Check memory usage again
+            new_memory_usage = self.redis.execute_command('MEMORY', 'USAGE', self.test_key)
+            assert new_memory_usage > 0, "MEMORY USAGE should still return a positive value after setting attributes"
+
+            # Memory usage should be higher after adding attributes
+            assert new_memory_usage > memory_usage, "Memory usage increase after adding attributes"
+
+        except Exception as e:
+            raise AssertionError(f"MEMORY USAGE command failed: {str(e)}")