From ba3cde6112ccc45810118215c661999d6851da24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Bonhomme?= <cedric@cedricbonhomme.org>
Date: Wed, 10 Apr 2019 15:09:07 +0200
Subject: [PATCH] Updated tests.

---
 tests/test_lsb.py    | 14 +++++++++++++-
 tests/test_lsbset.py | 15 +++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/tests/test_lsb.py b/tests/test_lsb.py
index d186aec..9529a20 100644
--- a/tests/test_lsb.py
+++ b/tests/test_lsb.py
@@ -132,7 +132,6 @@ class TestLSB(unittest.TestCase):
         messages_to_hide = ["a", "foo", "Hello World!", ":Python:"]
 
         for message in messages_to_hide:
-            message = "Hello World"
             outputBytes = io.BytesIO()
             with open("./tests/sample-files/20160505T130442.jpg", 'rb') as f:
                 bytes_image = lsb.hide(f, message)
@@ -143,6 +142,19 @@ class TestLSB(unittest.TestCase):
 
                 self.assertEqual(message, clear_message)
 
+    def test_with_location_of_image_as_argument(self):
+        messages_to_hide = ["Hello World!"]
+
+        for message in messages_to_hide:
+            outputBytes = io.BytesIO()
+            bytes_image = lsb.hide("./tests/sample-files/20160505T130442.jpg", message)
+            bytes_image.save(outputBytes, "PNG")
+            outputBytes.seek(0)
+
+            clear_message = lsb.reveal(outputBytes)
+
+            self.assertEqual(message, clear_message)
+
     def tearDown(self):
         try:
             os.unlink("./image.png")
diff --git a/tests/test_lsbset.py b/tests/test_lsbset.py
index 20861df..82fe754 100644
--- a/tests/test_lsbset.py
+++ b/tests/test_lsbset.py
@@ -25,6 +25,7 @@ __date__ = "$Date: 2016/04/13 $"
 __revision__ = "$Date: 2017/05/04 $"
 __license__ = "GPLv3"
 
+import io
 import os
 import unittest
 from unittest.mock import patch
@@ -103,6 +104,20 @@ class TestLSBSet(unittest.TestCase):
             secret = lsbset.hide("./tests/sample-files/Lenna-grayscale.png",
                                     message_to_hide, generators.eratosthenes())
 
+    def test_with_location_of_image_as_argument(self):
+        messages_to_hide = ["Hello World!"]
+
+        for message in messages_to_hide:
+            outputBytes = io.BytesIO()
+            bytes_image = lsbset.hide("./tests/sample-files/20160505T130442.jpg", message,
+                                      generators.identity())
+            bytes_image.save(outputBytes, "PNG")
+            outputBytes.seek(0)
+
+            clear_message = lsbset.reveal(outputBytes, generators.identity())
+
+            self.assertEqual(message, clear_message)
+
     def test_auto_convert_rgb(self):
         message_to_hide = "Hello World!"
         secret = lsbset.hide("./tests/sample-files/Lenna-grayscale.png",