diff --git a/firefly/__init__.py b/firefly/__init__.py
index 44f2246..3fc23f6 100644
--- a/firefly/__init__.py
+++ b/firefly/__init__.py
@@ -1,3 +1,9 @@
 # Firefly API client package
 from .client import FireflyClient
 from .exceptions import FireflyAPIError, FireflyAuthError
+
+__all__ = [
+    "FireflyClient",
+    "FireflyAPIError",
+    "FireflyAuthError",
+]
diff --git a/firefly/client.py b/firefly/client.py
index bc7b95d..d5f25f6 100644
--- a/firefly/client.py
+++ b/firefly/client.py
@@ -137,5 +137,5 @@ def generate_image(
                 contentClass=resp_json.get("contentClass"),
                 _response=resp,
             )
-        except (KeyError, IndexError, TypeError) as e:
+        except (KeyError, IndexError, TypeError):
             raise FireflyAPIError(f"Unexpected response format: {resp}")
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 12e07bc..8875f57 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,10 +1,10 @@
 import os
-import shutil
 import tempfile
 import pytest
 from typer.testing import CliRunner
 from unittest import mock
 from firefly.cli import app, mock_image
+import re
 
 runner = CliRunner()
 
@@ -48,7 +48,7 @@ def test_generate_download_image(monkeypatch):
         with open(os.path.join(os.path.dirname(__file__), "images", "cat-coding.png"), "rb") as f:
             expected = f.read()
         assert content == expected
-        assert f"Downloaded image (" in result.output
+        assert "Downloaded image (" in result.output
 
 @mock.patch("subprocess.run")
 def test_generate_show_images(mock_run):
@@ -190,6 +190,10 @@ def test_generate_with_all_new_options(monkeypatch):
     assert result.exit_code != 0
     assert "content_class must be either 'photo' or 'art'" in result.output
 
+def strip_ansi(text):
+    ansi_escape = re.compile(r'\x1b\[[0-9;]*[mGKHF]')
+    return ansi_escape.sub('', text)
+
 def test_generate_invalid_json_style(monkeypatch):
     result = runner.invoke(
         app,
@@ -203,7 +207,8 @@ def test_generate_invalid_json_style(monkeypatch):
         ]
     )
     assert result.exit_code == 2
-    assert "Invalid JSON for --style" in result.output
+    assert "Invalid JSON for --style" in strip_ansi(result.output)
+
 
 def test_generate_invalid_json_structure(monkeypatch):
     result = runner.invoke(
@@ -218,7 +223,7 @@ def test_generate_invalid_json_structure(monkeypatch):
         ]
     )
     assert result.exit_code == 2
-    assert "Invalid JSON for --structure" in result.output
+    assert "Invalid JSON for --structure" in strip_ansi(result.output)
 
 def test_generate_invalid_num_variations(monkeypatch):
     # Test too low
diff --git a/tests/test_models.py b/tests/test_models.py
index eceae87..4757126 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,4 +1,3 @@
-import pytest
 from firefly.models import FireflyImageSize, FireflyImage, FireflyImageOutput, FireflyImageResponse
 from unittest.mock import Mock
 
@@ -45,4 +44,4 @@ def test_firefly_image_response_content_class():
     img = FireflyImage(url="http://example.com/image.png")
     output = FireflyImageOutput(seed=2, image=img)
     resp = FireflyImageResponse(size=size, outputs=[output], contentClass="test-class")
-    assert resp.contentClass == "test-class" 
\ No newline at end of file
+    assert resp.contentClass == "test-class"
\ No newline at end of file