""" Test cases for deltalakeopsserver/logstats/log_stats_deltalake.py """ import os import sys import sqlite3 import tempfile import shutil import unittest from unittest.mock import patch, MagicMock, Mock from datetime import datetime # Add project root to path project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')) sys.path.insert(0, project_root) # Add deltalakeopsserver/logstats to path so objstore can be imported logstats_path = os.path.join(project_root, 'deltalakeopsserver', 'logstats') sys.path.insert(0, logstats_path) try: from deltalake import write_deltalake, DeltaTable DELTALAKE_AVAILABLE = True except ImportError: DELTALAKE_AVAILABLE = False write_deltalake = None DeltaTable = None from deltalakeopsserver.logstats.log_stats_deltalake import ( write_log_stats_to_deltalake, write_telemetry_to_deltalake, process_sqlite_to_deltalake, write_to_deltalake, app, ) from deltalakeopsserver.logstats.objstore import Local class TestWriteLogStatsToDeltaLake(unittest.TestCase): """Test write_log_stats_to_deltalake function""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.temp_db = os.path.join(self.temp_dir, "test.db") # Create a test SQLite database conn = sqlite3.connect(self.temp_db) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS log_stats ( component TEXT, level TEXT, position TEXT, msg_digest TEXT, count INTEGER, sample TEXT, sample_object_key TEXT ) """) cursor.execute(""" INSERT INTO log_stats (component, level, position, msg_digest, count, sample, sample_object_key) VALUES (?, ?, ?, ?, ?, ?, ?) """, ("tidb", "ERROR", "file.py:123", "digest123", 10, "sample error", "key123")) conn.commit() conn.close() def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_deltalake') def test_write_log_stats_to_deltalake(self, mock_write_deltalake): """Test writing log stats to Delta Lake""" local_path = f"file://{self.temp_dir}/deltalake" write_log_stats_to_deltalake( temp_file_path=self.temp_db, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_path=local_path, date_int=20250101, hour_int=12, sqlite_s3_path="s3://bucket/test.db" ) # Verify write_deltalake was called mock_write_deltalake.assert_called_once() call_args = mock_write_deltalake.call_args self.assertEqual(call_args[0][0], local_path) self.assertEqual(call_args[1]['partition_by'], ['date', 'org', 'cluster']) self.assertEqual(call_args[1]['mode'], 'append') # Verify the PyArrow table has correct data pa_table = call_args[0][1] self.assertEqual(len(pa_table), 1) @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_deltalake') def test_write_log_stats_empty_data(self, mock_write_deltalake): """Test writing empty log stats (should not write)""" # Create empty database empty_db = os.path.join(self.temp_dir, "empty.db") conn = sqlite3.connect(empty_db) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS log_stats ( component TEXT, level TEXT, position TEXT, msg_digest TEXT, count INTEGER, sample TEXT, sample_object_key TEXT ) """) conn.commit() conn.close() write_log_stats_to_deltalake( temp_file_path=empty_db, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_path="file:///tmp/deltalake", date_int=20250101, hour_int=12, sqlite_s3_path="s3://bucket/empty.db" ) # Should not call write_deltalake for empty data mock_write_deltalake.assert_not_called() @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_deltalake') def test_write_log_stats_position_split(self, mock_write_deltalake): """Test that position field is correctly split into file and line""" local_path = f"file://{self.temp_dir}/deltalake" write_log_stats_to_deltalake( temp_file_path=self.temp_db, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_path=local_path, date_int=20250101, hour_int=12, sqlite_s3_path="s3://bucket/test.db" ) # Verify the table has file and line columns pa_table = mock_write_deltalake.call_args[0][1] df = pa_table.to_pandas() self.assertIn("file", df.columns) self.assertIn("line", df.columns) self.assertEqual(df.iloc[0]["file"], "file.py") self.assertEqual(df.iloc[0]["line"], "123") @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_deltalake') def test_write_log_stats_position_without_colon_still_writes(self, mock_write_deltalake): """Test that position without ':' character still gets written (file=position, line='')""" # Create db with position that has no colon db_no_colon = os.path.join(self.temp_dir, "no_colon.db") conn = sqlite3.connect(db_no_colon) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS log_stats ( component TEXT, level TEXT, position TEXT, msg_digest TEXT, count INTEGER, sample TEXT, sample_object_key TEXT ) """) cursor.execute(""" INSERT INTO log_stats (component, level, position, msg_digest, count, sample, sample_object_key) VALUES (?, ?, ?, ?, ?, ?, ?) """, ("tidb", "WARN", "some_module", "digest456", 5, "sample", "key456")) conn.commit() conn.close() local_path = f"file://{self.temp_dir}/deltalake" write_log_stats_to_deltalake( temp_file_path=db_no_colon, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_path=local_path, date_int=20250101, hour_int=12, sqlite_s3_path="s3://bucket/no_colon.db" ) # Should still call write_deltalake mock_write_deltalake.assert_called_once() pa_table = mock_write_deltalake.call_args[0][1] df = pa_table.to_pandas() self.assertEqual(len(df), 1) self.assertEqual(df.iloc[0]["file"], "some_module") self.assertEqual(df.iloc[0]["line"], "") class TestWriteTelemetryToDeltaLake(unittest.TestCase): """Test write_telemetry_to_deltalake function""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.temp_db = os.path.join(self.temp_dir, "test.db") # Create a test SQLite database with telemetry_stats conn = sqlite3.connect(self.temp_db) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS telemetry_stats ( ts INTEGER, instance TEXT, telemetry_data TEXT ) """) cursor.execute(""" INSERT INTO telemetry_stats (ts, instance, telemetry_data) VALUES (?, ?, ?) """, (1234567890, "instance1", '{"key": "value"}')) conn.commit() conn.close() def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_deltalake') def test_write_telemetry_to_deltalake(self, mock_write_deltalake): """Test writing telemetry to Delta Lake""" local_path = f"file://{self.temp_dir}/deltalake" write_telemetry_to_deltalake( temp_file_path=self.temp_db, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_path=local_path, date_int=20250101, sqlite_s3_path="s3://bucket/test.db" ) # Verify write_deltalake was called mock_write_deltalake.assert_called_once() call_args = mock_write_deltalake.call_args self.assertEqual(call_args[0][0], local_path) self.assertEqual(call_args[1]['partition_by'], ['date', 'org', 'cluster']) self.assertEqual(call_args[1]['mode'], 'append') # Verify the PyArrow table has correct data pa_table = call_args[0][1] self.assertEqual(len(pa_table), 1) @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_deltalake') def test_write_telemetry_empty_data(self, mock_write_deltalake): """Test writing empty telemetry (should not write)""" # Create empty database empty_db = os.path.join(self.temp_dir, "empty.db") conn = sqlite3.connect(empty_db) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS telemetry_stats ( ts INTEGER, instance TEXT, telemetry_data TEXT ) """) conn.commit() conn.close() write_telemetry_to_deltalake( temp_file_path=empty_db, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_path="file:///tmp/deltalake", date_int=20250101, sqlite_s3_path="s3://bucket/empty.db" ) # Should not call write_deltalake for empty data mock_write_deltalake.assert_not_called() class TestProcessSqliteToDeltaLake(unittest.TestCase): """Test process_sqlite_to_deltalake function""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.temp_db = os.path.join(self.temp_dir, "2025010112.db") # Create a test SQLite database conn = sqlite3.connect(self.temp_db) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS log_stats ( component TEXT, level TEXT, position TEXT, msg_digest TEXT, count INTEGER, sample TEXT, sample_object_key TEXT ) """) cursor.execute(""" INSERT INTO log_stats (component, level, position, msg_digest, count, sample, sample_object_key) VALUES (?, ?, ?, ?, ?, ?, ?) """, ("tidb", "ERROR", "file.py:123", "digest123", 10, "sample error", "key123")) cursor.execute(""" CREATE TABLE IF NOT EXISTS telemetry_stats ( ts INTEGER, instance TEXT, telemetry_data TEXT ) """) cursor.execute(""" INSERT INTO telemetry_stats (ts, instance, telemetry_data) VALUES (?, ?, ?) """, (1234567890, "instance1", '{"key": "value"}')) conn.commit() conn.close() def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_log_stats_to_deltalake') @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_telemetry_to_deltalake') @patch('deltalakeopsserver.logstats.log_stats_deltalake.S3') def test_process_sqlite_to_deltalake_local(self, mock_s3_class, mock_telemetry, mock_log_stats): """Test processing SQLite file with local file system""" sqlite_path = "s3://test-bucket/2025010112.db" # Mock S3 client download - copy the test DB to temp location mock_s3 = MagicMock() def mock_download(bucket, key, filepath): # Copy our test DB to the temp file path import shutil shutil.copy(self.temp_db, filepath) return True mock_s3.download = Mock(side_effect=mock_download) mock_s3_class.return_value = mock_s3 process_sqlite_to_deltalake( vendor="aws", sqlite_s3_path=sqlite_path, org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_base_path="s3://bucket/deltalake/" ) # Verify both functions were called mock_log_stats.assert_called_once() mock_telemetry.assert_called_once() def test_process_sqlite_invalid_path(self): """Test processing with invalid path scheme""" with patch('deltalakeopsserver.logstats.log_stats_deltalake.logger') as mock_logger: process_sqlite_to_deltalake( vendor="aws", sqlite_s3_path="invalid://bucket/file.db", org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_base_path="s3://bucket/deltalake/" ) # Should log error and return early mock_logger.error.assert_called() def test_process_sqlite_invalid_filename(self): """Test processing with invalid filename format""" with patch('deltalakeopsserver.logstats.log_stats_deltalake.logger') as mock_logger: process_sqlite_to_deltalake( vendor="aws", sqlite_s3_path="s3://bucket/invalid.db", org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_base_path="s3://bucket/deltalake/" ) # Should log error and return early mock_logger.error.assert_called() def test_process_sqlite_invalid_vendor(self): """Test processing with invalid vendor""" with self.assertRaises(ValueError): process_sqlite_to_deltalake( vendor="invalid", sqlite_s3_path="s3://bucket/2025010112.db", org="test-org", project="test-project", cluster="test-cluster", deltalake_s3_base_path="s3://bucket/deltalake/" ) class TestWriteToDeltaLake(unittest.TestCase): """Test write_to_deltalake function""" @patch('deltalakeopsserver.logstats.log_stats_deltalake.process_sqlite_to_deltalake') def test_write_to_deltalake_aws(self, mock_process): """Test write_to_deltalake with AWS vendor""" write_to_deltalake( vendor="aws", org="test-org", project="test-project", cluster="test-cluster", bucket="test-bucket", sqlite_s3_path="s3://bucket/2025010112.db" ) mock_process.assert_called_once() call_args = mock_process.call_args[1] self.assertEqual(call_args['deltalake_s3_base_path'], "s3://test-bucket/deltalake/global/") @patch('deltalakeopsserver.logstats.log_stats_deltalake.process_sqlite_to_deltalake') def test_write_to_deltalake_gcp(self, mock_process): """Test write_to_deltalake with GCP vendor""" write_to_deltalake( vendor="gcp", org="test-org", project="test-project", cluster="test-cluster", bucket="test-bucket", sqlite_s3_path="gs://bucket/2025010112.db" ) mock_process.assert_called_once() call_args = mock_process.call_args[1] self.assertEqual(call_args['deltalake_s3_base_path'], "gs://test-bucket/deltalake/global/") @patch('deltalakeopsserver.logstats.log_stats_deltalake.process_sqlite_to_deltalake') def test_write_to_deltalake_azure(self, mock_process): """Test write_to_deltalake with Azure vendor""" write_to_deltalake( vendor="azure", org="test-org", project="test-project", cluster="test-cluster", bucket="test-bucket", sqlite_s3_path="abfs://bucket/2025010112.db" ) mock_process.assert_called_once() call_args = mock_process.call_args[1] self.assertEqual(call_args['deltalake_s3_base_path'], "abfs://test-bucket/deltalake/global/") def test_write_to_deltalake_invalid_vendor(self): """Test write_to_deltalake with invalid vendor""" with self.assertRaises(ValueError): write_to_deltalake( vendor="invalid", org="test-org", project="test-project", cluster="test-cluster", bucket="test-bucket", sqlite_s3_path="s3://bucket/2025010112.db" ) class TestFlaskRoutes(unittest.TestCase): """Test Flask routes""" def setUp(self): """Set up test client""" app.config['TESTING'] = True self.client = app.test_client() def test_health_check(self): """Test health check endpoint""" response = self.client.get('/health') self.assertEqual(response.status_code, 200) data = response.get_json() self.assertEqual(data['message'], 'healthy') @patch('deltalakeopsserver.logstats.log_stats_deltalake.write_to_deltalake') def test_write_log_state_success(self, mock_write): """Test write log stats endpoint with valid data""" data = { 'vendor': 'aws', 'org': 'test-org', 'project': 'test-project', 'cluster': 'test-cluster', 'bucket': 'test-bucket', 'sqlite_s3_path': 's3://bucket/2025010112.db' } response = self.client.post('/deltalake/log_stats', json=data) self.assertEqual(response.status_code, 202) result = response.get_json() self.assertEqual(result['message'], 'task accepted') mock_write.assert_called_once() def test_write_log_state_no_data(self): """Test write log stats endpoint without data""" # Send null JSON value - request.get_json() will return None response = self.client.post( '/deltalake/log_stats', data='null', content_type='application/json' ) # When request.get_json() returns None, if not data is True, returning 400 self.assertEqual(response.status_code, 400) result = response.get_json() self.assertEqual(result['error'], 'need data') def test_write_log_state_invalid_vendor(self): """Test write log stats endpoint with invalid vendor""" data = { 'vendor': 'invalid', 'org': 'test-org', 'project': 'test-project', 'cluster': 'test-cluster', 'bucket': 'test-bucket', 'sqlite_s3_path': 's3://bucket/2025010112.db' } response = self.client.post('/deltalake/log_stats', json=data) self.assertEqual(response.status_code, 400) result = response.get_json() self.assertEqual(result['error'], 'invalid vendor') if __name__ == '__main__': unittest.main()