From fe118fa685f2db6676c888a467bd0bae92ad01fa Mon Sep 17 00:00:00 2001
From: datalore <d4t4.l0r3@protonmail.ch>
Date: Thu, 3 Oct 2024 13:47:48 +0200
Subject: [PATCH] test(generatePrimes): Added benchmark functions

---
 main_test.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 main_test.go

diff --git a/main_test.go b/main_test.go
new file mode 100644
index 0000000..90147f3
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,21 @@
+package main
+
+import "testing"
+
+func testGeneratePrimes(upperLimit int64, useTensors bool, b *testing.B) {
+	var config Config
+	config.DontLoad = true
+	config.UseTensors = useTensors
+
+	for i := 0; i < b.N; i++ {
+		generatePrimes(upperLimit, config)
+	}
+}
+
+func BenchmarkPrimes1000CPU(b *testing.B) {
+	testGeneratePrimes(1000, false, b)
+}
+
+func BenchmarkPrimes1000GPU(b *testing.B) {
+	testGeneratePrimes(1000, true, b)
+}