scrape-giant/tests/test_browser_session.py

import sqlite3
import tempfile
import unittest
from pathlib import Path
from unittest import mock

import browser_session
import scrape_costco


class BrowserSessionTests(unittest.TestCase):
    def test_read_firefox_local_storage_reads_copied_sqlite(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            profile_dir = Path(tmpdir) / "abcd.default-release"
            ls_dir = profile_dir / "storage" / "default" / "https+++www.costco.com" / "ls"
            ls_dir.mkdir(parents=True)
            db_path = ls_dir / "data.sqlite"

            with sqlite3.connect(db_path) as connection:
                connection.execute("CREATE TABLE data (key TEXT, value TEXT)")
                connection.execute(
                    "INSERT INTO data (key, value) VALUES (?, ?)",
                    ("costco-x-wcs-clientId", "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf"),
                )

            values = browser_session.read_firefox_local_storage(
                profile_dir,
                origin_filter="costco.com",
            )

            self.assertEqual(
                "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf",
                values["costco-x-wcs-clientId"],
            )

    def test_load_costco_browser_headers_reads_id_token_and_client_id(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            profile_dir = Path(tmpdir)
            storage_dir = profile_dir / "storage" / "default" / "https+++www.costco.com" / "ls"
            storage_dir.mkdir(parents=True)
            db_path = storage_dir / "data.sqlite"

            with sqlite3.connect(db_path) as connection:
                connection.execute("CREATE TABLE data (key TEXT, value TEXT)")
                connection.execute(
                    "INSERT INTO data (key, value) VALUES (?, ?)",
                    ("idToken", "header.payload.signature"),
                )
                connection.execute(
                    "INSERT INTO data (key, value) VALUES (?, ?)",
                    ("clientID", "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf"),
                )

            headers = scrape_costco.load_costco_browser_headers(
                profile_dir,
                authorization="",
                client_id="",
                client_identifier="481b1aec-aa3b-454b-b81b-48187e28f205",
            )

        self.assertEqual("Bearer header.payload.signature", headers["costco-x-authorization"])
        self.assertEqual(
            "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf",
            headers["costco-x-wcs-clientId"],
        )
        self.assertEqual(
            "481b1aec-aa3b-454b-b81b-48187e28f205",
            headers["client-identifier"],
        )

    def test_load_costco_browser_headers_prefers_env_values(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            profile_dir = Path(tmpdir)
            storage_dir = profile_dir / "storage" / "default" / "https+++www.costco.com" / "ls"
            storage_dir.mkdir(parents=True)
            db_path = storage_dir / "data.sqlite"

            with sqlite3.connect(db_path) as connection:
                connection.execute("CREATE TABLE data (key TEXT, value TEXT)")
                connection.execute(
                    "INSERT INTO data (key, value) VALUES (?, ?)",
                    ("idToken", "storage.payload.signature"),
                )
                connection.execute(
                    "INSERT INTO data (key, value) VALUES (?, ?)",
                    ("clientID", "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf"),
                )

            headers = scrape_costco.load_costco_browser_headers(
                profile_dir,
                authorization="Bearer env.payload.signature",
                client_id="env-client-id",
                client_identifier="481b1aec-aa3b-454b-b81b-48187e28f205",
            )

        self.assertEqual("Bearer env.payload.signature", headers["costco-x-authorization"])
        self.assertEqual("env-client-id", headers["costco-x-wcs-clientId"])

    def test_scrape_costco_prompts_for_profile_dir_when_autodiscovery_fails(self):
        with mock.patch.object(
            scrape_costco,
            "find_firefox_profile_dir",
            side_effect=FileNotFoundError("no default profile"),
        ), mock.patch.object(
            scrape_costco.click,
            "prompt",
            return_value=Path("/tmp/profile"),
        ) as mocked_prompt, mock.patch.object(
            scrape_costco,
            "load_config",
            return_value={
                "authorization": "",
                "client_id": "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf",
                "client_identifier": "481b1aec-aa3b-454b-b81b-48187e28f205",
            },
        ), mock.patch.object(
            scrape_costco,
            "load_costco_browser_headers",
            return_value={
                "costco-x-authorization": "Bearer header.payload.signature",
                "costco-x-wcs-clientId": "4900eb1f-0c10-4bd9-99c3-c59e6c1ecebf",
                "client-identifier": "481b1aec-aa3b-454b-b81b-48187e28f205",
            },
        ), mock.patch.object(
            scrape_costco,
            "build_session",
            return_value=object(),
        ), mock.patch.object(
            scrape_costco,
            "fetch_summary_windows",
            return_value=(
                {"data": {"receiptsWithCounts": {"receipts": []}}},
                [],
            ),
        ), mock.patch.object(
            scrape_costco,
            "write_json",
        ), mock.patch.object(
            scrape_costco,
            "write_csv",
        ):
            scrape_costco.main.callback(
                outdir="/tmp/costco_output",
                document_type="all",
                document_sub_type="all",
                window_days=92,
                months_back=3,
                firefox_profile_dir=None,
            )

        mocked_prompt.assert_called_once()


if __name__ == "__main__":
    unittest.main()