stack/tests/conf/test_clients.py

"""Tests for conf._nessie, conf._polaris, conf._s3 client wrappers."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import httpx

from conf._nessie import NessieClient
from conf._polaris import PolarisClient
from conf._s3 import S3Client


class TestNessieClient:
    def test_init(self):
        c = NessieClient("http://localhost:19120/api/v2")
        assert c.base_url == "http://localhost:19120/api/v2"
        c.close()

    def test_config(self):
        c = NessieClient("http://localhost:19120/api/v2")
        with patch.object(c._client, "get") as mock:
            mock.return_value = MagicMock(json=lambda: {"maxSupportedApiVersion": 2})
            result = c.config()
            assert result["maxSupportedApiVersion"] == 2
        c.close()

    def test_list_refs(self):
        c = NessieClient("http://test")
        resp = MagicMock()
        resp.json.return_value = {"references": [{"name": "main", "type": "BRANCH"}]}
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            refs = c.list_refs()
            assert len(refs) == 1
            assert refs[0]["name"] == "main"
        c.close()

    def test_get_ref(self):
        c = NessieClient("http://test")
        resp = MagicMock()
        resp.json.return_value = {"name": "main", "hash": "abc123"}
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            ref = c.get_ref("main")
            assert ref["name"] == "main"
        c.close()

    def test_create_branch(self):
        c = NessieClient("http://test")
        get_resp = MagicMock()
        get_resp.json.return_value = {"name": "main", "hash": "abc123"}
        get_resp.raise_for_status = MagicMock()
        post_resp = MagicMock()
        post_resp.json.return_value = {"name": "dev", "hash": "abc123"}
        post_resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=get_resp):
            with patch.object(c._client, "post", return_value=post_resp):
                result = c.create_branch("dev")
                assert result["name"] == "dev"
        c.close()

    def test_list_contents(self):
        c = NessieClient("http://test")
        resp = MagicMock()
        resp.json.return_value = {"entries": [{"name": {"elements": ["ns", "t1"]}}]}
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            entries = c.list_contents()
            assert len(entries) == 1
        c.close()


class TestPolarisClient:
    def test_init_without_secret(self):
        c = PolarisClient("http://test:8181")
        assert c._token == ""
        c.close()

    def test_authenticate(self):
        c = PolarisClient.__new__(PolarisClient)
        c.base_url = "http://test"
        c.secret = "s3cr3t"
        c._client = httpx.Client(base_url="http://test", timeout=30.0)
        c._token = ""
        resp = MagicMock()
        resp.json.return_value = {"access_token": "tok123"}
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "post", return_value=resp):
            c._authenticate()
            assert c._token == "tok123"
            assert "Bearer tok123" in c._client.headers.get("Authorization", "")
        c.close()

    def test_init_with_secret_authenticates(self):
        resp = MagicMock()
        resp.json.return_value = {"access_token": "autotok"}
        resp.raise_for_status = MagicMock()
        with patch("httpx.Client.post", return_value=resp):
            c = PolarisClient("http://test", secret="mysecret")
            assert c._token == "autotok"
            c.close()

    def test_list_catalogs(self):
        c = PolarisClient("http://test")
        resp = MagicMock()
        resp.json.return_value = {"catalogs": [{"name": "aco"}]}
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            cats = c.list_catalogs()
            assert len(cats) == 1
            assert cats[0]["name"] == "aco"
        c.close()

    def test_list_namespaces(self):
        c = PolarisClient("http://test")
        resp = MagicMock()
        resp.json.return_value = {"namespaces": [["core"], ["pfs"]]}
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            ns = c.list_namespaces("aco")
            assert len(ns) == 2
        c.close()


class TestS3Client:
    def test_init(self):
        c = S3Client("http://test:9000", bucket="mybucket")
        assert c.bucket == "mybucket"
        c.close()

    def test_list_buckets(self):
        c = S3Client("http://test:9000", access_key="ak", secret_key="sk")
        xml = (
            '<?xml version="1.0"?>'
            "<ListAllMyBucketsResult>"
            "<Buckets><Bucket><Name>lakehouse</Name></Bucket>"
            "<Bucket><Name>exports</Name></Bucket></Buckets>"
            "</ListAllMyBucketsResult>"
        )
        resp = MagicMock(text=xml)
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            buckets = c.list_buckets()
            assert "lakehouse" in buckets
            assert "exports" in buckets
        c.close()

    def test_list_buckets_no_auth(self):
        c = S3Client("http://test:9000")
        xml = (
            '<?xml version="1.0"?>'
            "<ListAllMyBucketsResult><Buckets></Buckets></ListAllMyBucketsResult>"
        )
        resp = MagicMock(text=xml)
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            buckets = c.list_buckets()
            assert buckets == []
        c.close()

    def test_list_objects(self):
        c = S3Client("http://test:9000", access_key="ak", secret_key="sk")
        xml = (
            '<?xml version="1.0"?>'
            "<ListBucketResult>"
            "<Contents><Key>file1.parquet</Key></Contents>"
            "<Contents><Key>file2.parquet</Key></Contents>"
            "</ListBucketResult>"
        )
        resp = MagicMock(text=xml)
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            keys = c.list_objects(prefix="data/")
            assert "file1.parquet" in keys
            assert "file2.parquet" in keys
        c.close()

    def test_list_objects_no_auth(self):
        c = S3Client("http://test:9000")
        xml = '<?xml version="1.0"?><ListBucketResult></ListBucketResult>'
        resp = MagicMock(text=xml)
        resp.raise_for_status = MagicMock()
        with patch.object(c._client, "get", return_value=resp):
            keys = c.list_objects()
            assert keys == []
        c.close()