LihonT/tests/test_profile_flags.py

"""Native flag reads are routed through LihonTFlags, and LihonTFlags obeys the settings.

The smali rewriter is exercised on small synthetic methods. The runtime class is
compiled and run on a real JVM against the production settings classes; Android
itself is only on the classpath, as in test_settings_model.
"""

import os
from pathlib import Path
import shutil
import subprocess
import tempfile
import unittest

from lihont import profile_features

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src" / "lihont" / "two"
HARNESS = ROOT / "tests" / "java" / "lihont" / "two" / "FlagsTest.java"
HOOK = "Llihont/two/LihonTFlags;"

ALL_FLAGS = {**profile_features.PROFILE_FLAGS, **profile_features.REPOST_FLAGS}

INT_READ = """\
.class public Lcom/example/Reader;
.super Ljava/lang/Object;

.method public static read()I
    .registers 3

    const-string v0, "profile_left_align"

    const/4 v1, 0x0

    invoke-static {v0, v1}, Lcom/bytedance/ies/abmock/SettingsManager;->LIZ(Ljava/lang/String;I)I

    move-result v2

    return v2
.end method
"""

BOOL_READ = INT_READ.replace("(Ljava/lang/String;I)I", "(Ljava/lang/String;Z)Z").replace(
    "read()I", "read()Z")

PARAMETER_READ = """\
.class public Lcom/example/Reader;
.super Ljava/lang/Object;

.method public run(I)I
    .registers 5

    const-string v0, "repost_support_digg"

    invoke-static {v0, p1}, LX/0abc;->LIZ(Ljava/lang/String;I)I

    move-result p1

    return p1
.end method
"""

LEGACY = """\
.class public Lcom/example/Reader;
.super Ljava/lang/Object;

.method public static read()I
    .registers 3

    const-string v0, "profile_left_align"

    const/4 v1, 0x0

    invoke-static {v0, v1}, Lcom/bytedance/ies/abmock/SettingsManager;->LIZ(Ljava/lang/String;I)I

    move-result v2
    # lihont feature flag: profile_left_align
    const/16 v2, 0x1

    return v2
.end method
"""


class RewriteTest(unittest.TestCase):
    def test_int_read_goes_through_the_hook(self):
        text, hits = profile_features.rewrite_smali(INT_READ, ALL_FLAGS)
        self.assertEqual(hits, {"profile_left_align": 1})
        self.assertIn(
            "    move-result v2\n"
            "    # lihont feature flag: profile_left_align\n"
            "    invoke-static/range {v2 .. v2}, %s->profile_left_align(I)I\n"
            "    move-result v2\n" % HOOK, text)
        self.assertNotIn("const/16", text)
        self.assertIn("SettingsManager;->LIZ(Ljava/lang/String;I)I", text)
        self.assertIn("    return v2\n", text)

    def test_boolean_read_uses_the_boolean_overload(self):
        text, hits = profile_features.rewrite_smali(BOOL_READ, ALL_FLAGS)
        self.assertEqual(hits, {"profile_left_align": 1})
        self.assertIn("%s->profile_left_align(Z)Z" % HOOK, text)

    def test_parameter_register_is_kept(self):
        text, hits = profile_features.rewrite_smali(PARAMETER_READ, ALL_FLAGS)
        self.assertEqual(hits, {"repost_support_digg": 1})
        self.assertIn("invoke-static/range {p1 .. p1}, %s->repost_support_digg(I)I" % HOOK, text)
        self.assertIn("    move-result p1\n\n    return p1\n", text)

    def test_unknown_keys_are_left_alone(self):
        text, hits = profile_features.rewrite_smali(
            INT_READ.replace("profile_left_align", "some_other_flag"), ALL_FLAGS)
        self.assertEqual(hits, {})
        self.assertNotIn("LihonTFlags", text)

    def test_rewriting_twice_changes_nothing(self):
        once, _ = profile_features.rewrite_smali(INT_READ, ALL_FLAGS)
        twice, hits = profile_features.rewrite_smali(once, ALL_FLAGS)
        self.assertEqual(hits, {})
        self.assertEqual(once, twice)

    def test_old_const_override_is_replaced(self):
        text, hits = profile_features.rewrite_smali(LEGACY, ALL_FLAGS)
        self.assertEqual(hits, {"profile_left_align": 1})
        self.assertNotIn("const/16", text)
        self.assertEqual(text.count("# lihont feature flag"), 1)
        self.assertIn("(I)I\n    move-result v2\n\n    return v2\n", text)

    def test_marker_without_a_known_block_does_not_eat_code(self):
        stray = INT_READ.replace(
            "    move-result v2\n",
            "    move-result v2\n    # lihont feature flag: profile_left_align\n")
        text, hits = profile_features.rewrite_smali(stray, ALL_FLAGS)
        self.assertEqual(hits, {"profile_left_align": 1})
        self.assertIn("    return v2\n", text)


class RuntimeTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.javac = shutil.which("javac")
        cls.java = shutil.which("java")
        if not cls.javac or not cls.java:
            raise RuntimeError("Flag runtime tests require a JDK with java and javac")
        cls.android_jar = ROOT / "tools" / "android.jar"
        if not cls.android_jar.is_file():
            raise RuntimeError("tools/android.jar must be available")
        cls.temporary = tempfile.TemporaryDirectory(prefix="lihont-flags-tests-")
        cls.addClassCleanup(cls.temporary.cleanup)
        cls.classes = Path(cls.temporary.name) / "classes"
        cls.classes.mkdir()
        cls.environment = dict(os.environ)
        cls.environment.pop("JAVA_TOOL_OPTIONS", None)
        settings = SRC / "settings"
        cls.run_command([
            cls.javac, "--release", "17", "-encoding", "UTF-8",
            "-cp", str(cls.android_jar), "-d", str(cls.classes),
            str(settings / "ModSettings.java"), str(settings / "SettingsCodec.java"),
            str(settings / "SettingsRepository.java"), str(settings / "AndroidSettingsStorage.java"),
            str(settings / "SettingsHost.java"), str(SRC / "LihonTFlags.java"), str(HARNESS),
        ])

    @classmethod
    def run_command(cls, command):
        result = subprocess.run(
            command, capture_output=True, text=True, encoding="utf-8",
            env=cls.environment, timeout=60,
        )
        if result.returncode:
            raise AssertionError(
                f"Command failed with exit {result.returncode}: {command[0]}\n"
                + result.stdout + result.stderr
            )
        return result

    @staticmethod
    def entries():
        return (["banners:%s:%d" % item for item in profile_features.PROFILE_FLAGS.items()]
                + ["reposts:%s:%d" % item for item in profile_features.REPOST_FLAGS.items()])

    def run_scenario(self, name):
        result = self.run_command([
            self.java, "-cp", str(self.classes) + os.pathsep + str(self.android_jar),
            "lihont.two.FlagsTest", name,
        ] + self.entries())
        return result.stdout.strip()

    def test_each_switch_controls_only_its_own_flags(self):
        self.assertEqual(self.run_scenario("matrix"), "PASS matrix")

    def test_saved_change_applies_to_the_next_process(self):
        self.assertEqual(self.run_scenario("snapshot"), "PASS snapshot")

    def test_defaults_apply_while_settings_cannot_be_read(self):
        self.assertEqual(self.run_scenario("unavailable"), "PASS unavailable")

    def test_runtime_offers_exactly_what_the_patcher_calls(self):
        offered = set(self.run_scenario("surface").splitlines())
        called = {target.split("->", 1)[1] for target in profile_features.runtime_targets()}
        self.assertEqual(offered, called)


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