升级 NUPACK 4.1 并支持混合材料设计
This commit is contained in:
Lihatoo
parent
5daa60a464
commit
c6189d857d
33 changed files with 5830 additions and 466 deletions
525
service/test_server_calculations.py
Normal file
525
service/test_server_calculations.py
Normal file
|
|
@ -0,0 +1,525 @@
|
|||
import copy
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
os.environ.setdefault("NP_ACCOUNT_DB_PATH", "/tmp/np-replica-calculation-test.sqlite3")
|
||||
os.environ.setdefault("ENABLE_RNAPLOT", "0")
|
||||
|
||||
import server
|
||||
from nupack import Complex, Model, SetSpec, Strand, Tube, pfunc, tube_analysis
|
||||
|
||||
|
||||
def analysis_payload():
|
||||
payload = copy.deepcopy(server.EXAMPLE_PAYLOAD)
|
||||
payload["model"] = {
|
||||
"material": "rna",
|
||||
"ensemble": "stacking",
|
||||
"celsius": 37,
|
||||
"sodium": 1.0,
|
||||
"magnesium": 0.0,
|
||||
}
|
||||
payload["strands"] = [
|
||||
{"name": "A", "sequence": "ACG", "concentration": 1.0, "unit": "uM"},
|
||||
{"name": "B", "sequence": "CGU", "concentration": 0.5, "unit": "uM"},
|
||||
]
|
||||
payload["compute"] = ["pfunc", "pairs"]
|
||||
payload["options"].update(
|
||||
{
|
||||
"sparsity_fraction": 1.0,
|
||||
"sparsity_threshold": 0.0,
|
||||
"result_limit": 99,
|
||||
"pairs_preview_size": 24,
|
||||
}
|
||||
)
|
||||
payload["tube"] = {
|
||||
"name": "tube-test",
|
||||
"max_size": 2,
|
||||
"include_complexes": "A+B+B",
|
||||
"exclude_complexes": "A+A",
|
||||
}
|
||||
return payload
|
||||
|
||||
|
||||
class OfficialTubePipelineTest(unittest.TestCase):
|
||||
def test_replica_matches_direct_tube_analysis(self):
|
||||
payload = analysis_payload()
|
||||
replica = server.run_job_payload(payload)
|
||||
|
||||
model = Model(**payload["model"])
|
||||
a = Strand("ACG", name="A")
|
||||
b = Strand("CGU", name="B")
|
||||
included = Complex([a, b, b], name="included")
|
||||
excluded = Complex([a, a], name="excluded")
|
||||
tube = Tube(
|
||||
strands={a: 1e-6, b: 0.5e-6},
|
||||
complexes=SetSpec(max_size=2, include=[included], exclude=[excluded]),
|
||||
name="tube-test",
|
||||
)
|
||||
direct = tube_analysis(
|
||||
[tube],
|
||||
model=model,
|
||||
compute=["pairs"],
|
||||
options={
|
||||
"num_sample": 20,
|
||||
"energy_gap": 1.0,
|
||||
"sparsity_fraction": 1.0,
|
||||
"sparsity_threshold": 0.0,
|
||||
"single_mfe": False,
|
||||
},
|
||||
)
|
||||
|
||||
expected = {
|
||||
server.stringify_complex(complex_obj): float(value)
|
||||
for complex_obj, value in direct[tube].complex_concentrations.items()
|
||||
}
|
||||
actual = {
|
||||
row["display"]: row["concentration_M"]
|
||||
for row in replica["tube"]["complex_concentrations"]
|
||||
}
|
||||
|
||||
self.assertEqual(set(actual), set(expected))
|
||||
for name, value in expected.items():
|
||||
self.assertAlmostEqual(actual[name], value, places=18)
|
||||
|
||||
self.assertEqual(replica["tube"]["include_complexes"], ["A + B + B"])
|
||||
self.assertEqual(replica["tube"]["exclude_complexes"], ["A + A"])
|
||||
self.assertIsNotNone(replica["tube"]["ensemble_pair_fractions"])
|
||||
self.assertAlmostEqual(
|
||||
replica["tube"]["fraction_bases_unpaired"],
|
||||
float(direct[tube].fraction_bases_unpaired),
|
||||
places=12,
|
||||
)
|
||||
self.assertEqual(replica["performance"]["pipeline"], "tube_analysis")
|
||||
|
||||
def test_result_limit_is_display_only(self):
|
||||
payload = analysis_payload()
|
||||
payload["options"]["result_limit"] = 1
|
||||
replica = server.run_job_payload(payload)
|
||||
|
||||
self.assertGreater(replica["total_complex_count"], 1)
|
||||
self.assertEqual(replica["displayed_complex_count"], 1)
|
||||
self.assertEqual(len(replica["tube"]["complex_concentrations"]), 1)
|
||||
self.assertIsNotNone(replica["tube"]["ensemble_pair_fractions"])
|
||||
|
||||
def test_analysis_rejects_iupac_ambiguity_codes(self):
|
||||
payload = analysis_payload()
|
||||
payload["strands"][0]["sequence"] = "ACN"
|
||||
with self.assertRaisesRegex(ValueError, "only A, C, G, T, or U"):
|
||||
server.run_job_payload(payload)
|
||||
|
||||
|
||||
class OfficialUtilitiesPipelineTest(unittest.TestCase):
|
||||
def test_pfunc_matches_direct_utility(self):
|
||||
payload = copy.deepcopy(server.UTILITIES_EXAMPLE_PAYLOAD)
|
||||
payload["utility"]["operation"] = "pfunc"
|
||||
replica = server.run_job_payload(payload)
|
||||
direct_partition, direct_free_energy = pfunc(
|
||||
[payload["strands"][0]["sequence"]],
|
||||
Model(**payload["model"]),
|
||||
)
|
||||
self.assertEqual(replica["utility"]["partition_function"], str(direct_partition))
|
||||
self.assertAlmostEqual(
|
||||
replica["utility"]["free_energy_kcal_per_mol"],
|
||||
float(direct_free_energy),
|
||||
places=12,
|
||||
)
|
||||
|
||||
def test_distance_uses_two_explicit_inputs(self):
|
||||
payload = copy.deepcopy(server.UTILITIES_EXAMPLE_PAYLOAD)
|
||||
payload["strands"] = []
|
||||
payload["utility"].update(operation="seq_distance", input_a="AAAA", input_b="AAAU")
|
||||
replica = server.run_job_payload(payload)
|
||||
self.assertEqual(replica["utility"]["distance"], 1)
|
||||
|
||||
def test_pairs_matches_direct_utility(self):
|
||||
payload = copy.deepcopy(server.UTILITIES_EXAMPLE_PAYLOAD)
|
||||
payload["utility"]["operation"] = "pairs"
|
||||
payload["options"].update(sparsity_fraction=1.0, sparsity_threshold=0.0)
|
||||
replica = server.run_job_payload(payload)
|
||||
direct = server.nupack.pairs(
|
||||
[payload["strands"][0]["sequence"]],
|
||||
Model(**payload["model"]),
|
||||
sparsity_fraction=1.0,
|
||||
sparsity_threshold=0.0,
|
||||
)
|
||||
expected = server.serialize_pairs(direct, preview_limit=payload["options"]["pairs_preview_size"])
|
||||
self.assertEqual(replica["utility"]["pairs"]["shape"], expected["shape"])
|
||||
self.assertEqual(replica["utility"]["pairs"]["preview"], expected["preview"])
|
||||
|
||||
|
||||
class OfficialInputValidationTest(unittest.TestCase):
|
||||
def test_model_rejects_non_finite_and_invalid_physical_values(self):
|
||||
for field, value, message in (
|
||||
("celsius", float("nan"), "finite"),
|
||||
("celsius", -273.15, "absolute zero"),
|
||||
("sodium", float("inf"), "finite"),
|
||||
("sodium", -0.01, "non-negative"),
|
||||
("magnesium", -0.01, "non-negative"),
|
||||
):
|
||||
payload = analysis_payload()
|
||||
payload["model"][field] = value
|
||||
with self.subTest(field=field, value=value):
|
||||
with self.assertRaisesRegex(ValueError, message):
|
||||
server.run_job_payload(payload)
|
||||
|
||||
def test_probability_options_reject_values_outside_unit_interval(self):
|
||||
for field, value in (("sparsity_fraction", -0.1), ("sparsity_threshold", 1.01)):
|
||||
payload = analysis_payload()
|
||||
payload["options"][field] = value
|
||||
with self.subTest(field=field):
|
||||
with self.assertRaisesRegex(ValueError, "between 0 and 1"):
|
||||
server.run_job_payload(payload)
|
||||
|
||||
def test_analysis_concentration_must_be_finite(self):
|
||||
payload = analysis_payload()
|
||||
payload["strands"][0]["concentration"] = float("nan")
|
||||
with self.assertRaisesRegex(ValueError, "finite"):
|
||||
server.run_job_payload(payload)
|
||||
|
||||
|
||||
class Nupack41FeatureTest(unittest.TestCase):
|
||||
def test_material_specific_salt_rules_and_provenance(self):
|
||||
merna = server.parse_model_input({"material": "merna06", "celsius": 37})
|
||||
self.assertEqual(merna["sodium"], 0.12)
|
||||
self.assertEqual(merna["magnesium"], 0.0)
|
||||
with self.assertRaisesRegex(ValueError, "must be 0.12 M"):
|
||||
server.parse_model_input({"material": "merna06", "sodium": 1.0, "magnesium": 0})
|
||||
with self.assertRaisesRegex(ValueError, "between 0.12 and 1 M"):
|
||||
server.parse_model_input({"material": "rna-dna06", "sodium": 0.1, "magnesium": 0})
|
||||
|
||||
summary = server.build_model_summary({"material": "dna", "sodium": 1, "magnesium": 0})
|
||||
self.assertEqual(summary["resolved_material"], "dna04.3")
|
||||
self.assertEqual(summary["nupack_version"], "4.1.0.1")
|
||||
|
||||
def test_mixed_material_utility_preserves_prefixes_and_counts_bases(self):
|
||||
payload = copy.deepcopy(server.UTILITIES_EXAMPLE_PAYLOAD)
|
||||
payload["model"].update(material="rna-dna06", sodium=0.5, magnesium=0.0)
|
||||
payload["strands"] = [{"name": "hybrid", "sequence": "rACGdAT"}]
|
||||
payload["utility"]["operation"] = "pfunc"
|
||||
replica = server.run_job_payload(payload)
|
||||
|
||||
direct_partition, direct_free_energy = pfunc(
|
||||
["rACGdAT"], Model(material="rna-dna06", sodium=0.5, magnesium=0.0)
|
||||
)
|
||||
self.assertEqual(replica["strands"][0]["sequence"], "rACGdAT")
|
||||
self.assertEqual(replica["performance"]["workload"]["input_nucleotides"], 5)
|
||||
self.assertEqual(replica["utility"]["partition_function"], str(direct_partition))
|
||||
self.assertAlmostEqual(replica["utility"]["free_energy_kcal_per_mol"], float(direct_free_energy), places=12)
|
||||
|
||||
def test_mixed_material_complex_analysis_uses_true_nucleotide_lengths(self):
|
||||
payload = analysis_payload()
|
||||
payload["mode"] = "complex"
|
||||
payload["model"].update(material="rna-dna06", sodium=0.5, magnesium=0.0)
|
||||
payload["strands"] = [
|
||||
{"name": "H", "sequence": "rACGdAT", "concentration": 1.0, "unit": "uM"}
|
||||
]
|
||||
payload["complexes_text"] = "H"
|
||||
payload["compute"] = ["pfunc"]
|
||||
replica = server.run_job_payload(payload)
|
||||
|
||||
self.assertEqual(replica["performance"]["workload"]["input_nucleotides"], 5)
|
||||
self.assertEqual(replica["performance"]["workload"]["largest_complex_nucleotides"], 5)
|
||||
self.assertEqual(replica["complexes"][0]["strand_lengths"], [5])
|
||||
self.assertEqual(replica["model"]["resolved_material"], "rna-dna06")
|
||||
|
||||
def test_mixed_material_requires_explicit_lowercase_prefix(self):
|
||||
payload = copy.deepcopy(server.UTILITIES_EXAMPLE_PAYLOAD)
|
||||
payload["model"].update(material="rna-dna06", sodium=1.0, magnesium=0.0)
|
||||
payload["strands"] = [{"name": "hybrid", "sequence": "ACGdAT"}]
|
||||
payload["utility"]["operation"] = "pfunc"
|
||||
with self.assertRaisesRegex(ValueError, "explicit lowercase material prefixes"):
|
||||
server.run_job_payload(payload)
|
||||
|
||||
def test_mfe_new_controls_are_forwarded(self):
|
||||
payload = copy.deepcopy(server.UTILITIES_EXAMPLE_PAYLOAD)
|
||||
payload["utility"]["operation"] = "mfe"
|
||||
payload["options"].update(max_subopt_count=4321, indistinguishable_search=True)
|
||||
with mock.patch.object(server.nupack, "mfe", return_value=[]) as mfe:
|
||||
replica = server.run_job_payload(payload)
|
||||
mfe.assert_called_once_with(
|
||||
[payload["strands"][0]["sequence"]],
|
||||
mock.ANY,
|
||||
max_subopt_count=4321,
|
||||
indistinguishable_search=True,
|
||||
)
|
||||
self.assertEqual(replica["options"]["max_subopt_count"], 4321)
|
||||
self.assertTrue(replica["options"]["indistinguishable_search"])
|
||||
|
||||
def test_max_subopt_count_has_a_server_safety_bound(self):
|
||||
payload = analysis_payload()
|
||||
payload["options"]["max_subopt_count"] = 1000001
|
||||
with self.assertRaisesRegex(ValueError, "between 1 and 1000000"):
|
||||
server.run_job_payload(payload)
|
||||
|
||||
def test_analysis_complex_bonus_shifts_free_energy(self):
|
||||
payload = analysis_payload()
|
||||
payload["mode"] = "complex"
|
||||
payload["strands"] = [payload["strands"][0]]
|
||||
payload["compute"] = ["pfunc"]
|
||||
payload["complexes_text"] = "A"
|
||||
baseline = server.run_job_payload(payload)
|
||||
|
||||
payload["complexes_text"] = "A; bonus=1.25"
|
||||
shifted = server.run_job_payload(payload)
|
||||
self.assertAlmostEqual(
|
||||
shifted["complexes"][0]["free_energy_kcal_mol"]
|
||||
- baseline["complexes"][0]["free_energy_kcal_mol"],
|
||||
1.25,
|
||||
places=6,
|
||||
)
|
||||
self.assertEqual(shifted["complexes"][0]["bonus_kcal_mol"], 1.25)
|
||||
|
||||
def test_mixed_material_design_preserves_prefixes_and_true_lengths(self):
|
||||
payload = copy.deepcopy(server.DESIGN_TUBE_EXAMPLE_PAYLOAD)
|
||||
payload["workflow"] = "design"
|
||||
payload["mode"] = "complex"
|
||||
payload["model"].update(material="rna-dna06", sodium=0.5, magnesium=0.0)
|
||||
payload["design_domains"] = [
|
||||
{"name": "rseg", "sequence": "rN4"},
|
||||
{"name": "dseg", "sequence": "dN4"},
|
||||
]
|
||||
payload["strands"] = [
|
||||
{"name": "R", "sequence": "rseg"},
|
||||
{"name": "D", "sequence": "dseg"},
|
||||
]
|
||||
payload["design_complexes"] = [
|
||||
{"name": "RD", "strands": "R+D", "structure": "(4+)4"},
|
||||
]
|
||||
payload["design"].update(stop_condition=0.5, trials=1, seed=1)
|
||||
replica = server.run_job_payload(payload)
|
||||
|
||||
self.assertEqual(replica["model"]["material"], "rna-dna06")
|
||||
self.assertEqual([row["length"] for row in replica["design"]["domains"]], [4, 4])
|
||||
self.assertEqual([row["length"] for row in replica["strands"]], [4, 4])
|
||||
self.assertTrue(replica["strands"][0]["sequence"].startswith("r"))
|
||||
self.assertTrue(replica["strands"][1]["sequence"].startswith("d"))
|
||||
|
||||
|
||||
class OfficialDesignPipelineTest(unittest.TestCase):
|
||||
def test_tube_set_bonus_and_multi_tube_results(self):
|
||||
payload = copy.deepcopy(server.DESIGN_TUBE_EXAMPLE_PAYLOAD)
|
||||
payload["design_domains"] = [{"name": "a", "sequence": "N4"}]
|
||||
payload["design_complexes"][0].update(structure="(4+)4", bonus=0.5)
|
||||
payload["design"]["stop_condition"] = 0.5
|
||||
payload["design_tubes"] = [
|
||||
{
|
||||
"name": "T1",
|
||||
"max_size": 1,
|
||||
"include_complexes": "A+A",
|
||||
"exclude_complexes": "A",
|
||||
"on_targets": [{"complex": "AB_target", "concentration": 1, "unit": "uM"}],
|
||||
},
|
||||
{
|
||||
"name": "T2",
|
||||
"max_size": 1,
|
||||
"on_targets": [{"complex": "AB_target", "concentration": 2, "unit": "uM"}],
|
||||
},
|
||||
]
|
||||
replica = server.run_job_payload(payload)
|
||||
|
||||
self.assertEqual(replica["design"]["targets"][0]["bonus"], 0.5)
|
||||
self.assertEqual([tube["name"] for tube in replica["tubes"]], ["T1", "T2"])
|
||||
self.assertEqual(replica["tubes"][0]["include_complexes"], ["A+A"])
|
||||
self.assertEqual(replica["tubes"][0]["exclude_complexes"], ["A"])
|
||||
self.assertTrue(all(row["tube_name"] == "T1" for row in replica["tubes"][0]["complex_concentrations"]))
|
||||
self.assertTrue(all(row["tube_name"] == "T2" for row in replica["tubes"][1]["complex_concentrations"]))
|
||||
self.assertAlmostEqual(
|
||||
replica["design"]["objective"],
|
||||
replica["design"]["weighted_ensemble_defect"],
|
||||
places=12,
|
||||
)
|
||||
|
||||
|
||||
class SchedulingAndWorkloadTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
server.JOB_STORE.clear()
|
||||
server.JOB_DEDUP_RESERVATIONS.clear()
|
||||
|
||||
def tearDown(self):
|
||||
server.JOB_STORE.clear()
|
||||
server.JOB_DEDUP_RESERVATIONS.clear()
|
||||
|
||||
def test_cyclic_complex_count_for_four_species_through_size_four(self):
|
||||
self.assertEqual(server.cyclic_complex_count(4, 4), 108)
|
||||
|
||||
def test_design_workload_estimates_four_species(self):
|
||||
strands = [SimpleNamespace(name=name) for name in ("S1", "S2", "S3", "S4")]
|
||||
target = SimpleNamespace(strands=strands)
|
||||
summary = server.build_design_workload_summary(
|
||||
[{"name": "x", "constraint": "N10", "mutable": True}],
|
||||
[{"name": "target", "mutable": True}, {"name": "fixed", "mutable": False}],
|
||||
[{
|
||||
"name": "tube",
|
||||
"max_size": 4,
|
||||
"on_targets": [{"complex": "target"}],
|
||||
"include_complexes": [],
|
||||
}],
|
||||
{"target": target},
|
||||
{"max_time_seconds": 3600},
|
||||
)
|
||||
self.assertEqual(summary["mutable_domain_nucleotides"], 10)
|
||||
self.assertEqual(summary["fixed_target_count"], 1)
|
||||
self.assertEqual(summary["estimated_complexes_upper_bound"], 108)
|
||||
self.assertEqual(summary["estimated_off_targets_upper_bound"], 107)
|
||||
self.assertEqual(summary["largest_tube_complexes_upper_bound"], 108)
|
||||
|
||||
def test_design_workload_includes_excluded_species_and_explicit_large_complexes(self):
|
||||
a = SimpleNamespace(name="A")
|
||||
summary = server.build_design_workload_summary(
|
||||
[],
|
||||
[{"name": "target", "mutable": True}],
|
||||
[{
|
||||
"name": "tube",
|
||||
"max_size": 1,
|
||||
"on_targets": [{"complex": "target"}],
|
||||
"include_complexes": ["A+B+A"],
|
||||
"exclude_complexes": ["B"],
|
||||
}],
|
||||
{"target": SimpleNamespace(strands=[a])},
|
||||
{"max_time_seconds": 0},
|
||||
)
|
||||
# Two monomers are generated, B is excluded, and A+B+A is explicitly included.
|
||||
self.assertEqual(summary["estimated_complexes_upper_bound"], 2)
|
||||
self.assertEqual(summary["estimated_off_targets_upper_bound"], 1)
|
||||
|
||||
def test_cyclic_identity_treats_rotations_as_the_same_complex(self):
|
||||
self.assertEqual(
|
||||
server.canonical_cyclic_identity(("A", "B", "C")),
|
||||
server.canonical_cyclic_identity(("B", "C", "A")),
|
||||
)
|
||||
|
||||
def test_fixed_domain_complement_uses_material_alphabet(self):
|
||||
rna_model = Model(material="rna")
|
||||
dna_model = Model(material="dna")
|
||||
domain_map, domains = server.build_design_domains(
|
||||
[{"name": "a", "sequence": "AC"}], model=rna_model, material="rna"
|
||||
)
|
||||
_, rna_rows = server.build_design_strands(
|
||||
[{"name": "R", "sequence": "~a"}], domain_map, domains,
|
||||
material="rna", model=rna_model,
|
||||
)
|
||||
domain_map, domains = server.build_design_domains(
|
||||
[{"name": "a", "sequence": "AC"}], model=dna_model, material="dna"
|
||||
)
|
||||
_, dna_rows = server.build_design_strands(
|
||||
[{"name": "D", "sequence": "~a"}], domain_map, domains,
|
||||
material="dna", model=dna_model,
|
||||
)
|
||||
self.assertEqual(rna_rows[0]["fixed_sequence"], "GU")
|
||||
self.assertEqual(dna_rows[0]["fixed_sequence"], "GT")
|
||||
|
||||
def test_mixed_constraint_parser_keeps_material_segments(self):
|
||||
model = Model(material="rna-dna06", sodium=0.5)
|
||||
self.assertTrue(server.is_valid_iupac_constraint("rN4dA2wS2", "rna-dna06", model))
|
||||
self.assertEqual(server.expand_iupac_constraint("rN2dA2", "rna-dna06"), "rNNdAA")
|
||||
self.assertTrue(server.is_mutable_iupac_constraint("wA4", "rna-dna06"))
|
||||
self.assertFalse(server.is_valid_iupac_constraint("N4dN4", "rna-dna06", model))
|
||||
|
||||
def test_mixed_fixed_domain_complement_uses_material_alphabet(self):
|
||||
model = Model(material="rna-dna06", sodium=0.5)
|
||||
domain_map, domains = server.build_design_domains(
|
||||
[{"name": "x", "sequence": "rACGdAT"}], model=model, material="rna-dna06"
|
||||
)
|
||||
_, rows = server.build_design_strands(
|
||||
[{"name": "Xc", "sequence": "~x"}], domain_map, domains,
|
||||
material="rna-dna06", model=model,
|
||||
)
|
||||
self.assertEqual(rows[0]["fixed_sequence"], "dATrCGU")
|
||||
|
||||
def test_fingerprint_is_canonical_and_user_scoped(self):
|
||||
first = {"model": {"rna": True, "temperature": 37}, "compute": ["pairs", "mfe"]}
|
||||
reordered = {"compute": ["pairs", "mfe"], "model": {"temperature": 37, "rna": True}}
|
||||
self.assertEqual(
|
||||
server.job_payload_fingerprint(first, "alice"),
|
||||
server.job_payload_fingerprint(reordered, "alice"),
|
||||
)
|
||||
self.assertNotEqual(
|
||||
server.job_payload_fingerprint(first, "alice"),
|
||||
server.job_payload_fingerprint(first, "bob"),
|
||||
)
|
||||
|
||||
def test_active_duplicate_requires_force_and_is_scoped_by_user(self):
|
||||
owner = {"user_id": "alice"}
|
||||
other_owner = {"user_id": "bob"}
|
||||
payload = {"workflow": "analysis", "mode": "complex", "strands": []}
|
||||
with (
|
||||
mock.patch.object(server, "redis_enabled", return_value=False),
|
||||
mock.patch.object(server.ACCOUNT_STORE, "create_job"),
|
||||
mock.patch.object(server.threading, "Thread") as thread_class,
|
||||
):
|
||||
first_id, first_reused = server.create_job(payload, owner)
|
||||
duplicate_id, duplicate_reused = server.create_job(copy.deepcopy(payload), owner)
|
||||
self.assertFalse(first_reused)
|
||||
self.assertTrue(duplicate_reused)
|
||||
self.assertEqual(duplicate_id, first_id)
|
||||
self.assertEqual(thread_class.call_count, 1)
|
||||
|
||||
forced_id, forced_reused = server.create_job(payload, owner, force_duplicate=True)
|
||||
self.assertFalse(forced_reused)
|
||||
self.assertNotEqual(forced_id, first_id)
|
||||
self.assertEqual(thread_class.call_count, 2)
|
||||
|
||||
other_id, other_reused = server.create_job(payload, other_owner)
|
||||
self.assertFalse(other_reused)
|
||||
self.assertNotIn(other_id, {first_id, forced_id})
|
||||
self.assertEqual(thread_class.call_count, 3)
|
||||
|
||||
server.update_job_data(first_id, status="success", result={"ok": True}, payload=None)
|
||||
still_duplicate_id, still_duplicate = server.create_job(payload, owner)
|
||||
self.assertTrue(still_duplicate)
|
||||
self.assertEqual(still_duplicate_id, forced_id)
|
||||
|
||||
server.update_job_data(forced_id, status="success", result={"ok": True}, payload=None)
|
||||
next_id, next_reused = server.create_job(payload, owner)
|
||||
self.assertFalse(next_reused)
|
||||
self.assertNotIn(next_id, {first_id, forced_id})
|
||||
self.assertEqual(thread_class.call_count, 4)
|
||||
|
||||
def test_auto_resource_plan_respects_cpu_and_memory(self):
|
||||
with (
|
||||
mock.patch.object(server, "WORKER_CONCURRENCY", "auto"),
|
||||
mock.patch.object(server, "PER_JOB_THREAD_LIMIT", 4),
|
||||
mock.patch.object(server, "WORKER_MEMORY_RESERVE_GB", 4.0),
|
||||
mock.patch.object(server, "ESTIMATED_JOB_MEMORY_GB", 9.0),
|
||||
mock.patch.object(server, "detected_worker_cpu_count", return_value=64),
|
||||
mock.patch.object(server, "detected_worker_memory_gb", return_value=56.0),
|
||||
):
|
||||
plan = server.worker_resource_plan()
|
||||
self.assertEqual(plan["concurrency"], 5)
|
||||
self.assertEqual(plan["source"], "auto")
|
||||
self.assertEqual(plan["cpu_count"], 64)
|
||||
|
||||
def test_manual_resource_plan_keeps_explicit_concurrency(self):
|
||||
with (
|
||||
mock.patch.object(server, "WORKER_CONCURRENCY", "3"),
|
||||
mock.patch.object(server, "detected_worker_cpu_count", return_value=2),
|
||||
mock.patch.object(server, "detected_worker_memory_gb", return_value=8.0),
|
||||
):
|
||||
plan = server.worker_resource_plan()
|
||||
self.assertEqual(plan["concurrency"], 3)
|
||||
self.assertEqual(plan["source"], "manual")
|
||||
|
||||
def test_admin_session_is_random_revocable_and_does_not_expose_password(self):
|
||||
password = "admin-password-that-must-stay-secret"
|
||||
with (
|
||||
mock.patch.object(server, "redis_enabled", return_value=False),
|
||||
mock.patch.object(server, "ADMIN_TOKEN", password),
|
||||
mock.patch.object(server, "ADMIN_COOKIE_SECURE", False),
|
||||
):
|
||||
server.ADMIN_SESSION_STORE.clear()
|
||||
token = server.create_admin_session()
|
||||
self.assertNotEqual(token, password)
|
||||
self.assertTrue(server.valid_admin_session(token))
|
||||
cookie = server.admin_cookie_header(token)
|
||||
self.assertIn(f"{server.ADMIN_COOKIE_NAME}={token}", cookie)
|
||||
self.assertIn("HttpOnly", cookie)
|
||||
self.assertIn("SameSite=Strict", cookie)
|
||||
self.assertNotIn(password, cookie)
|
||||
server.delete_admin_session(token)
|
||||
self.assertFalse(server.valid_admin_session(token))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue