supaspin/tests/test_supaspin.py

"""Unit tests for the pure parts of supaspin. Run: python3 -m unittest discover -s tests"""
import base64
import hashlib
import hmac
import importlib.machinery
import importlib.util
import json
import os
import unittest

import yaml

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
loader = importlib.machinery.SourceFileLoader("supaspin", os.path.join(ROOT, "supaspin"))
spec = importlib.util.spec_from_loader("supaspin", loader)
sp = importlib.util.module_from_spec(spec)
loader.exec_module(sp)

CFG = {
    "storage_root": "/data/supabase", "node_name": "node1", "lan_ip": "192.168.10.30",
    "acme_email": "a@b.c", "ingress_class": "public",
    "slots": {1: "supa1.example.com", 2: "supa2.example.com", 3: "supa3.example.com"},
    "ssh_host": "203.0.113.5", "ssh_port": 2222, "ssh_user": "igor",
    "chart_version": "0.8.0",
}
MAIL = {"host": "smtp.x.com", "port": 587, "user": "u", "password": "p",
        "admin_email": "admin@x.com", "sender_name": "Sender"}


def template():
    with open(os.path.join(ROOT, "templates", "values.yaml.tpl")) as f:
        return f.read()


class Numbering(unittest.TestCase):
    def test_lowest_free_number_fills_gaps(self):
        self.assertEqual(sp.lowest_free_number(CFG["slots"], set()), 1)
        self.assertEqual(sp.lowest_free_number(CFG["slots"], {1, 3}), 2)
        self.assertEqual(sp.lowest_free_number(CFG["slots"], {1, 2}), 3)

    def test_no_free_slot(self):
        self.assertIsNone(sp.lowest_free_number(CFG["slots"], {1, 2, 3}))

    def test_ports_and_names(self):
        self.assertEqual((sp.pg_port(1), sp.pg_port(3)), (5441, 5443))
        self.assertEqual((sp.studio_port(1), sp.studio_port(3)), (3001, 3003))
        self.assertEqual(sp.namespace(2), "supaspin-2")
        self.assertEqual(sp.pvc_name(2, "db"), "supa2-supabase-db")
        self.assertEqual(sp.tls_secret(2), "supa2-tls")


class Jwt(unittest.TestCase):
    def test_signature_and_claims(self):
        tok = sp.sign_jwt("s" * 40, "anon", now=1000)
        h, p, s = tok.split(".")
        pad = lambda x: x + "=" * (-len(x) % 4)
        self.assertEqual(json.loads(base64.urlsafe_b64decode(pad(h))), {"alg": "HS256", "typ": "JWT"})
        claims = json.loads(base64.urlsafe_b64decode(pad(p)))
        self.assertEqual((claims["role"], claims["iss"], claims["iat"]), ("anon", "supabase", 1000))
        self.assertGreater(claims["exp"], 1000 + 9 * 365 * 86400)
        want = hmac.new(b"s" * 40, f"{h}.{p}".encode(), hashlib.sha256).digest()
        self.assertEqual(base64.urlsafe_b64decode(pad(s)), want)

    def test_generated_secrets_are_consistent(self):
        s = sp.generate_secrets()
        self.assertEqual(len(s["REALTIME_DB_ENC_KEY"]), 16)
        self.assertGreaterEqual(len(s["JWT_SECRET"]), 32)
        self.assertGreaterEqual(len(s["META_CRYPTO_KEY"]), 32)
        h, p, sig = s["ANON_KEY"].split(".")
        want = hmac.new(s["JWT_SECRET"].encode(), f"{h}.{p}".encode(), hashlib.sha256).digest()
        pad = lambda x: x + "=" * (-len(x) % 4)
        self.assertEqual(base64.urlsafe_b64decode(pad(sig)), want)
        self.assertNotEqual(sp.generate_secrets()["DB_PASSWORD"], s["DB_PASSWORD"])


class Values(unittest.TestCase):
    def setUp(self):
        self.sec = sp.generate_secrets()
        self.v = sp.instance_values(CFG, 2, self.sec, MAIL, template())

    def test_memory_caps_add_up(self):
        req, lim = sp.total_memory_mib()
        # db 2048, analytics 768, kong 512, realtime 384, the other eight 256 each
        self.assertEqual(lim, 2048 + 768 + 512 + 384 + 8 * 256)
        self.assertEqual(lim / 1024, 5.625)
        self.assertEqual(req, 512 + 11 * 128)  # 1.875 GiB
        # analytics and realtime were killed for running out of memory at half of production
        self.assertGreaterEqual(sp.MEMORY["analytics"][1], 768)
        self.assertGreaterEqual(sp.MEMORY["realtime"][1], 384)
        self.assertEqual(req / 1024, 1.875)
        for svc in sp.SERVICES:
            self.assertEqual(self.v["deployment"][svc]["priorityClassName"], "supaspin-low")
            self.assertIn("memory", self.v["deployment"][svc]["resources"]["limits"])

    def test_all_twelve_services_enabled(self):
        self.assertEqual(len(sp.SERVICES), 12)
        self.assertTrue(self.v["deployment"]["analytics"]["enabled"])
        self.assertTrue(self.v["deployment"]["vector"]["enabled"])

    def test_no_group_override_on_db_pod(self):
        # The db image's own user is 100:101; a forced fsGroup would lock Postgres out of its volume.
        self.assertNotIn("podSecurityContext", self.v["deployment"]["db"])
        self.assertEqual((sp.DB_UID, sp.DB_GID), (100, 101))

    def test_volume_caps(self):
        p = self.v["persistence"]
        self.assertEqual(p["db"]["size"], "2048Mi")
        self.assertEqual(p["storage"]["size"], "1024Mi")
        for vol in ("functions", "snippets", "deno", "pgsodium"):
            self.assertEqual(p[vol]["size"], "256Mi")
        self.assertEqual(set(sp.VOLUMES), {"db", "storage", "functions", "snippets", "deno", "pgsodium"})
        self.assertTrue(all(p[v]["storageClassName"] == "supaspin-local-sc" for v in sp.VOLUMES))

    def test_imgproxy_reads_the_storage_volume(self):
        img = self.v["persistence"]["imgproxy"]
        self.assertEqual(img["existingClaim"], "supa2-supabase-storage")
        self.assertNotIn("imgproxy", sp.VOLUMES)  # no separate image file for it

    def test_ingress_and_urls(self):
        ing = self.v["ingress"]
        self.assertEqual(ing["hosts"][0]["host"], "supa2.example.com")
        self.assertEqual(ing["tls"][0]["secretName"], "supa2-tls")
        self.assertEqual(ing["className"], "public")
        self.assertIsNone(ing["annotations"]["nginx.ingress.kubernetes.io/rewrite-target"])
        env = {e["name"]: e["value"] for e in self.v["environment"]["auth"]}
        self.assertEqual(env["API_EXTERNAL_URL"], "https://supa2.example.com")
        self.assertEqual(env["GOTRUE_SMTP_HOST"], "smtp.x.com")
        self.assertEqual(env["GOTRUE_SMTP_PORT"], "587")
        self.assertEqual(self.v["secret"]["smtp"]["password"], "p")

    def test_secrets_reach_values(self):
        self.assertEqual(self.v["secret"]["jwt"]["secret"], self.sec["JWT_SECRET"])
        self.assertEqual(self.v["secret"]["db"]["password"], self.sec["DB_PASSWORD"])

    def test_special_characters_stay_strings(self):
        m = dict(MAIL, password='p"a$s: #x', sender_name="A & B")
        v = sp.instance_values(CFG, 1, self.sec, m, template())
        self.assertEqual(v["secret"]["smtp"]["password"], 'p"a$s: #x')


class EnvFile(unittest.TestCase):
    def setUp(self):
        self.sec = sp.generate_secrets()
        self.text = sp.render_env(CFG, 1, self.sec, MAIL)
        self.kv = dict(l.split("=", 1) for l in self.text.splitlines() if "=" in l and not l.startswith("#"))

    def test_expected_variables(self):
        for k in ("SUPABASE_URL", "SUPABASE_PUBLIC_URL", "PUBLIC_SUPABASE_URL", "API_EXTERNAL_URL",
                  "PUBLIC_SUPABASE_ANON_KEY", "ANON_KEY", "JWT_ANON_KEY", "SERVICE_ROLE_KEY",
                  "JWT_SERVICE_KEY", "JWT_SECRET", "JWT_EXPIRY", "POSTGRES_PASSWORD", "POSTGRES_DB",
                  "POSTGRES_HOST", "POSTGRES_PORT", "DATABASE_URL", "DASHBOARD_USERNAME",
                  "DASHBOARD_PASSWORD", "GOTRUE_SMTP_HOST", "GOTRUE_SMTP_PORT", "GOTRUE_SMTP_USER",
                  "GOTRUE_SMTP_PASS", "GOTRUE_SMTP_ADMIN_EMAIL", "GOTRUE_SMTP_SENDER_NAME"):
            self.assertIn(k, self.kv)
        self.assertEqual(self.kv["SUPABASE_URL"], "https://supa1.example.com")

    def test_database_url_and_tunnel_comment(self):
        self.assertEqual(self.kv["DATABASE_URL"],
                         f"postgresql://postgres:{self.sec['DB_PASSWORD']}@localhost:5441/postgres")
        self.assertIn("192.168.10.30:5441", self.text)
        self.assertIn("ssh -N -p 2222 -L 5441:192.168.10.30:5441 -L 3001:192.168.10.30:3001 igor@203.0.113.5",
                      self.text)

    def test_mask_hides_secrets_only(self):
        masked = sp.mask_env(self.text)
        self.assertNotIn(self.sec["DB_PASSWORD"], masked)
        self.assertNotIn(self.sec["JWT_SECRET"], masked)
        self.assertIn("SUPABASE_URL=https://supa1.example.com", masked)
        self.assertIn("POSTGRES_PORT=5441", masked)


class Fstab(unittest.TestCase):
    def test_line_and_removal(self):
        line = sp.fstab_line("/r/1/db/db.img", "/r/1/db/data", "/r")
        self.assertEqual(line, "/r/1/db/db.img /r/1/db/data ext4 loop,nofail,x-systemd.requires-mounts-for=/r 0 0")
        text = "UUID=x / ext4 defaults 0 1\n" + line + "\n# " + line + "\n"
        self.assertTrue(sp.fstab_has(text, "/r/1/db/data"))
        out = sp.fstab_without(text, "/r/1/db/data")
        self.assertFalse(sp.fstab_has(out, "/r/1/db/data"))
        self.assertIn("UUID=x / ext4", out)
        self.assertIn("# /r/1/db/db.img", out)  # comments are left alone

    def test_mounts_under_an_instance_folder(self):
        text = ("a /r/1/db/data ext4 x 0 0\n"
                "b /r/1/imgproxy/data ext4 x 0 0\n"   # a volume an older version created
                "c /r/10/db/data ext4 x 0 0\n"        # another instance
                "# d /r/1/old/data ext4 x 0 0\n"
                "e /r/1 ext4 x 0 0\n")                # the folder itself is not below itself
        self.assertEqual(sp.fstab_mounts_under(text, "/r/1"), ["/r/1/db/data", "/r/1/imgproxy/data"])

    def test_similar_paths_are_not_removed(self):
        text = "a /r/1/db/data ext4 x 0 0\nb /r/10/db/data ext4 x 0 0\n"
        self.assertIn("/r/10/db/data", sp.fstab_without(text, "/r/1/db/data"))


RENDERED = """
apiVersion: apps/v1
kind: StatefulSet
metadata: {name: supa1-supabase-db}
spec: {template: {spec: {containers: [{name: db, ports: [{name: http, containerPort: 5432}]}]}}}
---
apiVersion: apps/v1
kind: Deployment
metadata: {name: supa1-supabase-studio}
spec: {template: {spec: {containers: [{name: studio, ports: [{name: http, containerPort: 3000}]}]}}}
---
apiVersion: apps/v1
kind: Deployment
metadata: {name: supa1-supabase-rest}
spec: {template: {spec: {containers: [{name: rest, ports: [{name: http, containerPort: 3000}]}]}}}
---
apiVersion: v1
kind: ConfigMap
metadata: {name: supa1-supabase-kong}
data:
  kong-entrypoint.sh: "echo hi\\necho there\\n"
  kong.yml: |
    _format_version: '2.1'
    services:
      - name: auth-v1
        url: http://auth:9999
        routes: [{name: auth-v1-all, paths: [/auth/v1/]}]
        plugins:
          - name: request-transformer
            config:
              add:
                headers:
                  - "Authorization: ${LUA_AUTH_EXPR}"
      - name: meta
        url: http://meta:8080/
        routes: [{name: meta-all, paths: [/pg/]}]
      - name: dashboard
        url: http://studio:3000/
        routes: [{name: dashboard-all, paths: [/]}]
"""


class PostRender(unittest.TestCase):
    def setUp(self):
        self.docs = list(yaml.safe_load_all(sp.post_render(RENDERED, 1, "192.168.10.30")))

    def find(self, kind, name):
        return next(d for d in self.docs if d["kind"] == kind and d["metadata"]["name"] == name)

    def test_db_and_studio_get_fixed_host_ports(self):
        db = self.find("StatefulSet", "supa1-supabase-db")["spec"]["template"]["spec"]["containers"][0]["ports"][0]
        self.assertEqual((db["hostPort"], db["hostIP"], db["containerPort"]), (5441, "192.168.10.30", 5432))
        st = self.find("Deployment", "supa1-supabase-studio")["spec"]["template"]["spec"]["containers"][0]["ports"][0]
        self.assertEqual((st["hostPort"], st["hostIP"]), (3001, "192.168.10.30"))

    def test_other_pods_untouched(self):
        rest = self.find("Deployment", "supa1-supabase-rest")["spec"]["template"]["spec"]["containers"][0]["ports"][0]
        self.assertNotIn("hostPort", rest)

    def test_studio_and_pgmeta_routes_removed(self):
        kong = yaml.safe_load(self.find("ConfigMap", "supa1-supabase-kong")["data"]["kong.yml"])
        self.assertEqual([s["name"] for s in kong["services"]], ["auth-v1"])
        self.assertEqual(kong["_format_version"], "2.1")

    def test_kong_file_is_cut_not_rewritten(self):
        # Kong's start script substitutes ${LUA_AUTH_EXPR} inside these exact double quotes.
        text = self.find("ConfigMap", "supa1-supabase-kong")["data"]["kong.yml"]
        self.assertIn('- "Authorization: ${LUA_AUTH_EXPR}"', text)
        self.assertIn("_format_version: '2.1'", text)
        self.assertNotIn("name: meta", text)
        self.assertNotIn("name: dashboard", text)

    def test_strip_keeps_following_top_level_keys(self):
        src = "services:\n  - name: meta\n    url: x\n  - name: keep\n    url: y\nother: 1\n"
        self.assertEqual(sp.strip_kong_services(src, ("meta",)),
                         "services:\n  - name: keep\n    url: y\nother: 1\n")
        src2 = "services:\n  - name: dashboard\n    url: x\nother: 1\n"
        self.assertEqual(sp.strip_kong_services(src2, ("dashboard",)), "services:\nother: 1\n")

    def test_other_config_keys_survive(self):
        cm = self.find("ConfigMap", "supa1-supabase-kong")
        self.assertEqual(cm["data"]["kong-entrypoint.sh"], "echo hi\necho there\n")

    def test_instance_number_selects_ports(self):
        docs = list(yaml.safe_load_all(sp.post_render(RENDERED.replace("supa1", "supa3"), 3, "10.0.0.1")))
        db = next(d for d in docs if d["kind"] == "StatefulSet")["spec"]["template"]["spec"]["containers"][0]["ports"][0]
        self.assertEqual((db["hostPort"], db["hostIP"]), (5443, "10.0.0.1"))


class Manifests(unittest.TestCase):
    def test_pv_is_pinned_and_reserved_for_the_chart_claim(self):
        pv = sp.pv_manifest(CFG, 2, "db")
        self.assertEqual(pv["spec"]["claimRef"], {"namespace": "supaspin-2", "name": "supa2-supabase-db"})
        self.assertEqual(pv["spec"]["local"]["path"], "/data/supabase/2/db/data")
        self.assertEqual(pv["spec"]["capacity"]["storage"], "2048Mi")
        self.assertEqual(pv["spec"]["persistentVolumeReclaimPolicy"], "Retain")
        self.assertEqual(pv["spec"]["nodeAffinity"]["required"]["nodeSelectorTerms"][0]
                         ["matchExpressions"][0]["values"], ["node1"])
        self.assertEqual(sp.pv_manifest(CFG, 2, "deno")["spec"]["capacity"]["storage"], "256Mi")

    def test_saved_certificate_keeps_issuer_annotations(self):
        src = {"type": "kubernetes.io/tls",
               "metadata": {"name": "supa1-tls", "namespace": "supaspin-1", "uid": "x",
                            "resourceVersion": "9", "ownerReferences": [{"name": "o"}],
                            "annotations": {"cert-manager.io/issuer-name": "supaspin-letsencrypt",
                                            "cert-manager.io/issuer-kind": "ClusterIssuer",
                                            "cert-manager.io/issuer-group": "cert-manager.io",
                                            "kubectl.kubernetes.io/last-applied-configuration": "{}"}},
               "data": {"tls.crt": "Yg==", "tls.key": "Yw=="}}
        out = sp.carry_secret(src, "supaspin-certs", "supa1-tls")
        self.assertEqual(out["metadata"], {"name": "supa1-tls", "namespace": "supaspin-certs",
                                           "annotations": {"cert-manager.io/issuer-name": "supaspin-letsencrypt",
                                                           "cert-manager.io/issuer-kind": "ClusterIssuer",
                                                           "cert-manager.io/issuer-group": "cert-manager.io"}})
        self.assertEqual(out["data"], src["data"])
        self.assertEqual(out["type"], "kubernetes.io/tls")

    def test_secret_without_annotations_still_copies(self):
        out = sp.carry_secret({"data": {"tls.crt": "a"}}, "ns", "n")
        self.assertNotIn("annotations", out["metadata"])
        self.assertEqual(out["type"], "kubernetes.io/tls")

    def test_tunnel_help_mentions_all_steps(self):
        text = sp.tunnel_help(CFG, 1, "PW")
        self.assertIn("postgresql://postgres:PW@localhost:5441/postgres", text)
        self.assertIn("http://localhost:3001", text)
        self.assertIn("192.168.10.30:5441", text)


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