From 1da1276d44002d8063d2c48a749035b98d4a4f3b Mon Sep 17 00:00:00 2001
From: Jen Hamon <jhamon@pinecone.io>
Date: Fri, 20 Sep 2024 19:22:12 -0400
Subject: [PATCH] Initial perf test setup

---
 poetry.lock                   | 33 ++++++++++++++++++++++++++++++++-
 pyproject.toml                |  1 +
 tests/perf/__init__.py        |  0
 tests/perf/docker-compose.yml | 20 ++++++++++++++++++++
 tests/perf/test_upsert.py     | 22 ++++++++++++++++++++++
 5 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 tests/perf/__init__.py
 create mode 100644 tests/perf/docker-compose.yml
 create mode 100644 tests/perf/test_upsert.py

diff --git a/poetry.lock b/poetry.lock
index 10f3864c..7367b704 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -977,6 +977,17 @@ files = [
 googleapis-common-protos = "*"
 protobuf = ">=4.21.0"
 
+[[package]]
+name = "py-cpuinfo"
+version = "9.0.0"
+description = "Get CPU info with pure Python"
+optional = false
+python-versions = "*"
+files = [
+    {file = "py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690"},
+    {file = "py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5"},
+]
+
 [[package]]
 name = "pygments"
 version = "2.16.1"
@@ -1030,6 +1041,26 @@ pytest = ">=5.4.0"
 [package.extras]
 testing = ["coverage", "hypothesis (>=5.7.1)"]
 
+[[package]]
+name = "pytest-benchmark"
+version = "4.0.0"
+description = "A ``pytest`` fixture for benchmarking code. It will group the tests into rounds that are calibrated to the chosen timer."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "pytest-benchmark-4.0.0.tar.gz", hash = "sha256:fb0785b83efe599a6a956361c0691ae1dbb5318018561af10f3e915caa0048d1"},
+    {file = "pytest_benchmark-4.0.0-py3-none-any.whl", hash = "sha256:fdb7db64e31c8b277dff9850d2a2556d8b60bcb0ea6524e36e28ffd7c87f71d6"},
+]
+
+[package.dependencies]
+py-cpuinfo = "*"
+pytest = ">=3.8"
+
+[package.extras]
+aspect = ["aspectlib"]
+elasticsearch = ["elasticsearch"]
+histogram = ["pygal", "pygaljs"]
+
 [[package]]
 name = "pytest-cov"
 version = "2.10.1"
@@ -1383,4 +1414,4 @@ grpc = ["googleapis-common-protos", "grpcio", "grpcio", "lz4", "protobuf", "prot
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8"
-content-hash = "7fcaf8376166ea98dd704c1ee4a27e09c64bfb2c598e4be8547ba3127a762f82"
+content-hash = "6055f6484a7e0a2bdcc39bf5d056cb7036d9a9108b1033d38198640264c5f745"
diff --git a/pyproject.toml b/pyproject.toml
index 12591369..b41cafd1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -101,6 +101,7 @@ pytest-timeout = "2.2.0"
 urllib3_mock = "0.3.3"
 responses = ">=0.8.1"
 black = "^24.4.2"
+pytest-benchmark = "^4.0.0"
 
 [tool.poetry.extras]
 grpc = ["grpcio", "googleapis-common-protos", "lz4", "protobuf", "protoc-gen-openapiv2"]
diff --git a/tests/perf/__init__.py b/tests/perf/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/perf/docker-compose.yml b/tests/perf/docker-compose.yml
new file mode 100644
index 00000000..250fd4c0
--- /dev/null
+++ b/tests/perf/docker-compose.yml
@@ -0,0 +1,20 @@
+services:
+  index1:
+    image: ghcr.io/pinecone-io/pinecone-index:latest
+    environment:
+      PORT: 5081
+      INDEX_TYPE: serverless
+      DIMENSION: 1536
+      METRIC: cosine
+    ports:
+      - "5081:5081"
+
+  index2:
+    image: ghcr.io/pinecone-io/pinecone-index:latest
+    environment:
+      PORT: 5082
+      INDEX_TYPE: pod
+      DIMENSION: 1536
+      METRIC: dot-product
+    ports:
+      - "5082:5082"
diff --git a/tests/perf/test_upsert.py b/tests/perf/test_upsert.py
new file mode 100644
index 00000000..3093a783
--- /dev/null
+++ b/tests/perf/test_upsert.py
@@ -0,0 +1,22 @@
+import random
+import uuid
+from pinecone.grpc import PineconeGRPC, GRPCClientConfig
+
+# Initialize a client. An API key must be passed, but the 
+# value does not matter.
+pc = PineconeGRPC(api_key="test_api_key")
+
+# Target the indexes. Use the host and port number along with disabling tls.
+index = pc.Index(host="localhost:5081", grpc_config=GRPCClientConfig(secure=False))
+dimension = 3
+
+def upserts():
+    vectors = []
+    for batch in range(0, 10):
+        for i in range(0, 100):
+            vectors.append((str(uuid.uuid4()), [random.random()] * dimension))
+        index.upsert(vectors=vectors, namespace="ns2")
+
+def test_upsert(benchmark):
+    benchmark(upserts)
+    print(index.describe_index_stats())