Fix linting issues in test files

Author: ipv1337

test_dir/test_main_improved.py

@@ -5,8 +5,8 @@
 
 import os
 import sys
-import unittest
-from unittest.mock import patch, MagicMock, call
+from unittest import TestCase, mock
+from unittest.mock import MagicMock, call, patch
 
 # Import pytest only if available
 try:
@@ -35,8 +35,12 @@ def decorator(f):
 try:
     from cli_code.config import Config
     from cli_code.main import (
-        cli, set_default_provider, set_default_model, list_models,
-        setup_command, start_interactive_session
+        cli,
+        list_models,
+        set_default_model,
+        set_default_provider,
+        setup_command,
+        start_interactive_session,
     )
     # Import click testing if available
     try:
@@ -69,7 +73,7 @@ def start_interactive_session(): pass
 )
 
 @skip_if_imports_unavailable
-class TestMainCLICommands(unittest.TestCase):
+class TestMainCLICommands(TestCase):
     """Tests for main CLI commands that weren't well covered previously."""
 
     @patch('cli_code.main.Config')
@@ -110,7 +114,7 @@ def test_list_models_invalid_provider(self, mock_get_models, mock_config):
         mock_get_models.assert_not_called()
 
 @skip_if_imports_unavailable
-class TestInteractiveSession(unittest.TestCase):
+class TestInteractiveSession(TestCase):
     """Tests for interactive session functionality."""
 
     @patch('cli_code.main.get_agent_for_provider')

test_dir/test_model_basic.py

@@ -3,8 +3,8 @@
 These tests focus on increasing coverage for the model classes.
 """
 
-import unittest
-from unittest.mock import patch, MagicMock
+from unittest import TestCase, skipIf
+from unittest.mock import MagicMock, patch
 
 # Import necessary modules safely
 try:
@@ -19,9 +19,10 @@ class AbstractModelAgent: pass
     class GeminiModelAgent: pass
     class OllamaModelAgent: pass
 
+
 # Skip all tests if imports aren't available
-@unittest.skipIf(not IMPORTS_AVAILABLE, "Required model imports not available")
-class TestGeminiModelBasics(unittest.TestCase):
+@skipIf(not IMPORTS_AVAILABLE, "Required model imports not available")
+class TestGeminiModelBasics(TestCase):
     """Test basic GeminiModelAgent functionality that doesn't require API calls."""
 
     def test_gemini_init(self):
@@ -82,8 +83,9 @@ def test_gemini_append_history(self, mock_get_model, mock_configure):
         self.assertEqual(agent.history[1]["role"], "model")
         self.assertEqual(agent.history[1]["parts"][0]["text"], "Hi there!")
 
-@unittest.skipIf(not IMPORTS_AVAILABLE, "Required model imports not available")
-class TestOllamaModelBasics(unittest.TestCase):
+
+@skipIf(not IMPORTS_AVAILABLE, "Required model imports not available")
+class TestOllamaModelBasics(TestCase):
     """Test basic OllamaModelAgent functionality that doesn't require API calls."""
 
     def test_ollama_init(self):