text: import glob
  import gzip
  import os
  import pickle
  import shutil
  from collections.abc import Iterator
  from datetime import datetime, timedelta
  from fnmatch import fnmatch
  from io import BytesIO
  from unittest import TestCase
  
  from seleniumwire.request import Request, Response, WebSocketMessage
  from seleniumwire.storage import RequestStorage
  
  
  class RequestStorageTest(TestCase):
  
      def test_initialise(self):
          RequestStorage(base_dir=self.base_dir)
          storage_dir = glob.glob(os.path.join(self.base_dir, '.seleniumwire', 'storage-*'))
  
          self.assertEqual(1, len(storage_dir))
  
      def test_cleanup_removes_storage(self):
          storage = RequestStorage(base_dir=self.base_dir)
          storage.cleanup()
  
          # The 'seleniumwire' parent folder should have been cleaned up
          # when there is nothing left inside of it.
          self.assertFalse(os.listdir(self.base_dir))
  
      def test_cleanup_does_not_remove_parent_folder(self):
          # There is an existing storage folder
          os.makedirs(os.path.join(self.base_dir, '.seleniumwire', 'teststorage'))
          storage = RequestStorage(base_dir=self.base_dir)
          storage.cleanup()
  
          # The existing storage folder is not cleaned up
          self.assertEqual(1, len(os.listdir(self.base_dir)))
          self.assertTrue(os.path.exists(os.path.join(self.base_dir, '.seleniumwire', 'teststorage')))
  
      def test_initialise_clears_old_folders(self):
          old_dir = os.path.join(self.base_dir, '.seleniumwire', 'storage-test1')
          new_dir = os.path.join(self.base_dir, '.seleniumwire', 'storage-test2')
          os.makedirs(old_dir)
          os.makedirs(new_dir)
          two_days_ago = (datetime.now() - timedelta(days=2)).timestamp()
          os.utime(old_dir, times=(two_days_ago, two_days_ago))
  
          RequestStorage(base_dir=self.base_dir)
  
          self.assertFalse(os.path.exists(old_dir))
          self.assertTrue(os.path.exists(new_dir))
  
      def test_save_request(self):
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
  
          storage.save_request(request)
  
          request_file_path = self._get_stored_path(request.id, 'request')
  
          with open(request_file_path[0], 'rb') as loaded:
              loaded_request = pickle.load(loaded)
  
          self.assertEqual(request.id, loaded_request.id)
          self.assertEqual('http://www.example.com/test/path/', loaded_request.url)
          self.assertEqual('GET', loaded_request.method)
          self.assertEqual({
              'Host': 'www.example.com',
              'Accept': '*/*'
          }, dict(loaded_request.headers))
          self.assertIsNone(loaded_request.response)
  
      def test_save_request_with_body(self):
          body = b'test request body'
          request = self._create_request(body=body)
          storage = RequestStorage(base_dir=self.base_dir)
  
          storage.save_request(request)
  
          request_file_path = self._get_stored_path(request.id, 'request')
  
          with open(request_file_path[0], 'rb') as loaded:
              loaded_request = pickle.load(loaded)
  
          self.assertEqual(body, loaded_request.body)
  
      def test_save_response(self):
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          response = self._create_response()
  
          storage.save_response(request.id, response)
  
          response_file_path = self._get_stored_path(request.id, 'response')
  
          with open(response_file_path[0], 'rb') as loaded:
              loaded_response = pickle.load(loaded)
  
          self.assertEqual(200, loaded_response.status_code)
          self.assertEqual('OK', loaded_response.reason)
          self.assertEqual({
              'Content-Type': 'application/json',
              'Content-Length': '500'
          }, dict(loaded_response.headers))
  
      def test_save_response_with_body(self):
          body = b'some response body'
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          response = self._create_response(body=body)
  
          storage.save_response(request.id, response)
  
          response_file_path = self._get_stored_path(request.id, 'response')
  
          with open(response_file_path[0], 'rb') as loaded:
              loaded_response = pickle.load(loaded)
  
          self.assertEqual(b'some response body', loaded_response.body)
  
      def test_save_response_no_request(self):
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          response = self._create_response()
          storage.clear_requests()
  
          storage.save_response(request.id, response)
  
          response_file_path = self._get_stored_path(request.id, 'response')
  
          self.assertFalse(response_file_path)
  
      def test_save_har_entry(self):
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
  
          storage.save_har_entry(request.id, {'name': 'test_har_entry'})
  
          har_file_path = self._get_stored_path(request.id, 'har_entry')
  
          with open(har_file_path[0], 'rb') as loaded:
              loaded_har = pickle.load(loaded)
  
          self.assertEqual(loaded_har['name'], 'test_har_entry')
  
      def test_save_har_entry_no_request(self):
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          storage.clear_requests()
  
          storage.save_har_entry(request.id, {'name': 'test_har_entry'})
  
          har_file_path = self._get_stored_path(request.id, 'har_entry')
  
          self.assertFalse(har_file_path)
  
      def test_load_requests(self):
          request_1 = self._create_request()
          request_2 = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request_1)
          storage.save_request(request_2)
  
          requests = storage.load_requests()
  
          self.assertEqual(2, len(requests))
          self.assertEqual(request_1.id, requests[0].id)
          self.assertEqual(request_2.id, requests[1].id)
          self.assertIsNone(requests[0].response)
          self.assertIsNone(requests[1].response)
  
      def test_iter_requests(self):
          request_1 = self._create_request()
          request_2 = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request_1)
          storage.save_request(request_2)
  
          requests = storage.iter_requests()
  
          self.assertIsInstance(requests, Iterator)
          requests = list(requests)
          self.assertEqual(2, len(requests))
          self.assertEqual(request_1.id, requests[0].id)
          self.assertEqual(request_2.id, requests[1].id)
  
      def test_load_response(self):
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          response = self._create_response()
          storage.save_response(request.id, response)
  
          requests = storage.load_requests()
  
          self.assertIsNotNone(requests[0].response)
  
      def test_load_response_encoded_body(self):
          body = b'test response body'
          io = BytesIO()
          with gzip.GzipFile(fileobj=io, mode='wb') as f:
              f.write(body)
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          response = self._create_response(body=io.getvalue())
          response.headers['Content-Encoding'] = 'gzip'
          storage.save_response(request.id, response)
  
          response_body = storage.load_requests()[0].response.body
  
          self.assertEqual(body, response_body)
  
      def test_load_response_encoded_body_error(self):
          body = b'test response body'
          request = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request)
          response = self._create_response(body=body)
          response.headers['Content-Encoding'] = 'gzip'
          storage.save_response(request.id, response)
  
          response_body = storage.load_requests()[0].response.body
  
          self.assertEqual(body, response_body)
  
      def test_load_last_request(self):
          request_1 = self._create_request()
          request_2 = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request_1)
          storage.save_request(request_2)
  
          last_request = storage.load_last_request()
  
          self.assertEqual(request_2.id, last_request.id)
  
      def test_load_last_request_none(self):
          storage = RequestStorage(base_dir=self.base_dir)
  
          last_request = storage.load_last_request()
  
          self.assertIsNone(last_request)
  
      def test_load_request_with_ws_messages(self):
          storage = RequestStorage(base_dir=self.base_dir)
          request_1 = self._create_request()  # Websocket handshake request
          request_2 = self._create_request()
          storage.save_request(request_1)
          storage.save_request(request_2)
          storage.save_ws_message(request_1.id, WebSocketMessage(
              from_client=True,
              content='websocket test message',
              date=datetime.now(),
          ))
  
          requests = storage.load_requests()
  
          self.assertTrue(len(requests[0].ws_messages) > 0)
          self.assertEqual('websocket test message', requests[0].ws_messages[0].content)
          self.assertTrue(len(requests[1].ws_messages) == 0)
  
      def test_load_request_cert_data(self):
          storage = RequestStorage(base_dir=self.base_dir)
          request = self._create_request()
          storage.save_request(request)
          response = self._create_response()
          response.cert = {'subject': 'test_cert'}
          storage.save_response(request.id, response)
  
          requests = storage.load_requests()
  
          self.assertEqual({'subject': 'test_cert'}, requests[0].cert)
  
      def test_clear_requests(self):
          request_1 = self._create_request()
          request_2 = self._create_request()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request_1)
          storage.save_request(request_2)
  
          storage.clear_requests()
          requests = storage.load_requests()
  
          self.assertFalse(requests)
          self.assertFalse(glob.glob(os.path.join(self.base_dir, '.seleniumwire', 'storage-*', '*')))
  
      def test_get_home_dir(self):
          storage = RequestStorage(base_dir=self.base_dir)
  
          self.assertEqual(os.path.join(self.base_dir, '.seleniumwire'), storage.home_dir)
  
      def test_get_session_dir(self):
          storage = RequestStorage(base_dir=self.base_dir)
  
          self.assertTrue(fnmatch(storage.session_dir, os.path.join(self.base_dir, '.seleniumwire', 'storage-*')))
  
      def test_find(self):
          request_1 = self._create_request('http://www.example.com/test/path/?foo=bar')
          request_2 = self._create_request('http://www.stackoverflow.com/other/path/?x=y')
          mock_response = self._create_response()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request_1)
          storage.save_response(request_1.id, mock_response)
          storage.save_request(request_2)
  
          self.assertEqual(request_1.id, storage.find('/test/path/').id)
          self.assertEqual(request_1.id, storage.find(r'/test/path/\?foo=bar').id)
          self.assertEqual(request_1.id, storage.find(r'http://www.example.com/test/path/\?foo=bar').id)
          self.assertEqual(request_1.id, storage.find(r'http://www.example.com/test/path/').id)
  
          self.assertIsNone(storage.find('/different/path'))
          self.assertIsNone(storage.find('/test/path/?x=y'))
          self.assertIsNone(storage.find(r'http://www.example.com/different/path/\?foo=bar'))
          self.assertIsNone(storage.find(r'http://www.different.com/test/path/\?foo=bar'))
          self.assertIsNone(storage.find(r'http://www.example.com/test/path/\?x=y'))
  
      def test_find_similar_urls(self):
          request_1 = self._create_request('https://192.168.1.1/redfish/v1')
          request_2 = self._create_request('https://192.168.1.1/redfish')
          mock_response = self._create_response()
          storage = RequestStorage(base_dir=self.base_dir)
          storage.save_request(request_1)
          storage.save_response(request_1.id, mock_response)
          storage.save_request(request_2)
          storage.save_response(request_2.id, mock_response)
  
          self.assertEqual(request_1.id, storage.find('.*v1').id)
          self.assertEqual(request_2.id, storage.find('https://192.168.1.1/redfish$').id)
  
      def _get_stored_path(self, request_id, filename):
          return glob.glob(os.path.join(self.base_dir, '.seleniumwire', 'storage-*',
                                        'request-{}'.format(request_id), filename))
  
      def _create_request(self, url='http://www.example.com/test/path/', body=b''):
          headers = [
              ('Host', 'www.example.com'),
              ('Accept', '*/*')
          ]
          return Request(method='GET', url=url, headers=headers, body=body)
  
      def _create_response(self, body=b''):
          headers = [
              ('Content-Type', 'application/json'),
              ('Content-Length', '500')
          ]
          return Response(status_code=200, reason='OK', headers=headers, body=body)
  
      def setUp(self):
          self.base_dir = os.path.join(os.path.dirname(__file__), 'data')
  
      def tearDown(self):
          shutil.rmtree(os.path.join(self.base_dir), ignore_errors=True)
  