import time
import sqlite3
import urllib.request
import unittest

DB_PATH = "orbital_forensics.db"
METRICS_URL = "http://localhost:9101/metrics"

class TestTelemetryPipeline(unittest.TestCase):

    def setUp(self):
        self.conn = sqlite3.connect(DB_PATH)
        cursor = self.conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS itar_file_transfers (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                file_name TEXT,
                encryption_status TEXT,
                ground_station_id TEXT,
                status TEXT
            )
        ''')
        cursor.execute('''
            INSERT INTO itar_file_transfers (file_name, encryption_status, ground_station_id, status)
            VALUES ('starlink_blueprint_v4.enc', 'Plaintext', 'GS-US-WEST-1', 'PENDING')
        ''')
        self.conn.commit()

    def test_prometheus_metric_increment(self):
        time.sleep(3)
        req = urllib.request.urlopen(METRICS_URL)
        metrics_output = req.read().decode('utf-8')

        self.assertIn("orbital_weak_encryption_streams_total", metrics_output)
        self.assertIn('file_name="starlink_blueprint_v4.enc"', metrics_output)
        self.assertIn('encryption_status="Plaintext"', metrics_output)

    def tearDown(self):
        cursor = self.conn.cursor()
        cursor.execute("DELETE FROM itar_file_transfers WHERE file_name = 'starlink_blueprint_v4.enc'")
        self.conn.commit()
        self.conn.close()

if __name__ == "__main__":
    unittest.main()