From ea476d73c6906da444e07004f5e9fcd0caee691c Mon Sep 17 00:00:00 2001
From: brianreicher <reicher.b@northeastern.edu>
Date: Tue, 5 Dec 2023 13:37:31 -0500
Subject: [PATCH] Db tests

---
 .github/workflows/tests.yaml | 11 ++++++
 tests/eval/db_test.py        | 72 ++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 tests/eval/db_test.py

diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index 575f4a6..d29c7b3 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -38,5 +38,16 @@ jobs:
       - name: Install dependencies
         run: pip install -r dev_requirments.txt
 
+      - name: Set up SQLite and create test database
+        run: |
+          sudo apt-get install sqlite3
+          python -m sqlite3 test_db.db <<EOF
+          CREATE TABLE test_scores_table (
+            network TEXT,
+            checkpoint INTEGER,
+            threshold REAL,
+            scores TEXT
+          );
+          EOF
       - name: Test with pytest
         run: pytest tests
\ No newline at end of file
diff --git a/tests/eval/db_test.py b/tests/eval/db_test.py
new file mode 100644
index 0000000..1cea50e
--- /dev/null
+++ b/tests/eval/db_test.py
@@ -0,0 +1,72 @@
+import unittest
+import os
+from autoseg.eval import Database
+
+
+class TestDatabase(unittest.TestCase):
+    def setUp(self):
+        # Create a test database
+        self.db_name = "test_db"
+        self.table_name = "test_scores_table"
+        self.db = Database(self.db_name, self.table_name)
+
+    def tearDown(self):
+        # Remove the test database file after tests
+        if os.path.exists(f"{self.db_name}.db"):
+            os.remove(f"{self.db_name}.db")
+
+    def test_add_score(self):
+        # Test adding a score to the database
+        network = "TestNetwork"
+        checkpoint = 1
+        threshold = 0.5
+        scores_dict = {"metric": 0.8}
+        self.db.add_score(network, checkpoint, threshold, scores_dict)
+
+        # Verify that the score is added correctly
+        result = self.db.get_scores(networks=network)
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0][0], network)
+        self.assertEqual(result[0][1], checkpoint)
+        self.assertEqual(result[0][2], threshold)
+        self.assertEqual(result[0][3], scores_dict)
+
+    def test_get_scores(self):
+        # Test retrieving scores from the database
+        network = "TestNetwork"
+        checkpoint = 1
+        threshold = 0.5
+        scores_dict = {"metric": 0.8}
+        self.db.add_score(network, checkpoint, threshold, scores_dict)
+
+        # Verify that the retrieved score matches the added score
+        result = self.db.get_scores(networks=network)
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0][0], network)
+        self.assertEqual(result[0][1], checkpoint)
+        self.assertEqual(result[0][2], threshold)
+        self.assertEqual(result[0][3], scores_dict)
+
+    def test_get_scores_multiple_conditions(self):
+        # Test retrieving scores with multiple conditions
+        network = "TestNetwork"
+        checkpoint = 1
+        threshold = 0.5
+        scores_dict = {"metric": 0.8}
+        self.db.add_score(network, checkpoint, threshold, scores_dict)
+
+        # Add another score with different conditions
+        network2 = "TestNetwork2"
+        checkpoint2 = 2
+        threshold2 = 0.7
+        scores_dict2 = {"metric": 0.9}
+        self.db.add_score(network2, checkpoint2, threshold2, scores_dict2)
+
+        # Verify that the retrieved scores match the added scores with multiple conditions
+        result = self.db.get_scores(networks=[network, network2], checkpoints=[checkpoint2])
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0][0], network2)
+        self.assertEqual(result[0][1], checkpoint2)
+        self.assertEqual(result[0][2], threshold2)
+        self.assertEqual(result[0][3], scores_dict2)
+