Skip to content

Commit 6614ca3

Browse files
committed
Updated/added unit tests accordingly
1 parent 8f84bb0 commit 6614ca3

File tree

5 files changed

+311
-2
lines changed

5 files changed

+311
-2
lines changed

src/models/config.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -587,7 +587,9 @@ class QuotaHandlersConfiguration(ConfigurationBase):
587587
sqlite: Optional[SQLiteDatabaseConfiguration] = None
588588
postgres: Optional[PostgreSQLDatabaseConfiguration] = None
589589
limiters: list[QuotaLimiterConfiguration] = Field(default_factory=list)
590-
scheduler: QuotaSchedulerConfiguration = Field(default_factory=QuotaSchedulerConfiguration)
590+
scheduler: QuotaSchedulerConfiguration = Field(
591+
default_factory=QuotaSchedulerConfiguration
592+
)
591593
enable_token_history: bool = False
592594

593595

@@ -610,7 +612,9 @@ class Configuration(ConfigurationBase):
610612
default_factory=ConversationCacheConfiguration
611613
)
612614
byok_rag: list[ByokRag] = Field(default_factory=list)
613-
quota_handlers: QuotaHandlersConfiguration = Field(default_factory=QuotaHandlersConfiguration)
615+
quota_handlers: QuotaHandlersConfiguration = Field(
616+
default_factory=QuotaHandlersConfiguration
617+
)
614618

615619
def dump(self, filename: str = "configuration.json") -> None:
616620
"""Dump actual configuration into JSON file."""

tests/unit/models/config/test_dump_configuration.py

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@
1313
PostgreSQLDatabaseConfiguration,
1414
CORSConfiguration,
1515
Configuration,
16+
QuotaHandlersConfiguration,
17+
QuotaLimiterConfiguration,
18+
QuotaSchedulerConfiguration,
1619
ServiceConfiguration,
1720
InferenceConfiguration,
1821
TLSConfiguration,
@@ -175,6 +178,8 @@ def test_dump_configuration(tmp_path) -> None:
175178
"quota_handlers": {
176179
"sqlite": None,
177180
"postgres": None,
181+
"limiters": [],
182+
"scheduler": {"period": 1},
178183
"enable_token_history": False,
179184
},
180185
}
@@ -293,3 +298,201 @@ def test_dump_configuration_with_more_mcp_servers(tmp_path) -> None:
293298
"url": "http://localhost:8083",
294299
},
295300
]
301+
302+
303+
def test_dump_configuration_with_quota_limiters(tmp_path) -> None:
304+
"""
305+
Test that the Configuration object can be serialized to a JSON file and
306+
that the resulting file contains all expected sections and values.
307+
308+
Please note that redaction process is not in place.
309+
"""
310+
cfg = Configuration(
311+
name="test_name",
312+
service=ServiceConfiguration(
313+
tls_config=TLSConfiguration(
314+
tls_certificate_path=Path("tests/configuration/server.crt"),
315+
tls_key_path=Path("tests/configuration/server.key"),
316+
tls_key_password=Path("tests/configuration/password"),
317+
),
318+
cors=CORSConfiguration(
319+
allow_origins=["foo_origin", "bar_origin", "baz_origin"],
320+
allow_credentials=False,
321+
allow_methods=["foo_method", "bar_method", "baz_method"],
322+
allow_headers=["foo_header", "bar_header", "baz_header"],
323+
),
324+
),
325+
llama_stack=LlamaStackConfiguration(
326+
use_as_library_client=True,
327+
library_client_config_path="tests/configuration/run.yaml",
328+
api_key="whatever",
329+
),
330+
user_data_collection=UserDataCollection(
331+
feedback_enabled=False, feedback_storage=None
332+
),
333+
database=DatabaseConfiguration(
334+
sqlite=None,
335+
postgres=PostgreSQLDatabaseConfiguration(
336+
db="lightspeed_stack",
337+
user="ls_user",
338+
password="ls_password",
339+
port=5432,
340+
ca_cert_path=None,
341+
ssl_mode="require",
342+
gss_encmode="disable",
343+
),
344+
),
345+
mcp_servers=[],
346+
customization=None,
347+
inference=InferenceConfiguration(
348+
default_provider="default_provider",
349+
default_model="default_model",
350+
),
351+
quota_handlers=QuotaHandlersConfiguration(
352+
limiters=[
353+
QuotaLimiterConfiguration(
354+
type="user_limiter",
355+
name="user_monthly_limits",
356+
initial_quota=1,
357+
quota_increase=10,
358+
period="2 seconds",
359+
),
360+
QuotaLimiterConfiguration(
361+
type="cluster_limiter",
362+
name="cluster_monthly_limits",
363+
initial_quota=2,
364+
quota_increase=20,
365+
period="1 month",
366+
),
367+
],
368+
scheduler=QuotaSchedulerConfiguration(period=10),
369+
enable_token_history=True,
370+
),
371+
)
372+
assert cfg is not None
373+
dump_file = tmp_path / "test.json"
374+
cfg.dump(dump_file)
375+
376+
with open(dump_file, "r", encoding="utf-8") as fin:
377+
content = json.load(fin)
378+
# content should be loaded
379+
assert content is not None
380+
381+
# all sections must exists
382+
assert "name" in content
383+
assert "service" in content
384+
assert "llama_stack" in content
385+
assert "user_data_collection" in content
386+
assert "mcp_servers" in content
387+
assert "authentication" in content
388+
assert "authorization" in content
389+
assert "customization" in content
390+
assert "inference" in content
391+
assert "database" in content
392+
assert "byok_rag" in content
393+
assert "quota_handlers" in content
394+
395+
# check the whole deserialized JSON file content
396+
assert content == {
397+
"name": "test_name",
398+
"service": {
399+
"host": "localhost",
400+
"port": 8080,
401+
"auth_enabled": False,
402+
"workers": 1,
403+
"color_log": True,
404+
"access_log": True,
405+
"tls_config": {
406+
"tls_certificate_path": "tests/configuration/server.crt",
407+
"tls_key_password": "tests/configuration/password",
408+
"tls_key_path": "tests/configuration/server.key",
409+
},
410+
"cors": {
411+
"allow_credentials": False,
412+
"allow_headers": [
413+
"foo_header",
414+
"bar_header",
415+
"baz_header",
416+
],
417+
"allow_methods": [
418+
"foo_method",
419+
"bar_method",
420+
"baz_method",
421+
],
422+
"allow_origins": [
423+
"foo_origin",
424+
"bar_origin",
425+
"baz_origin",
426+
],
427+
},
428+
},
429+
"llama_stack": {
430+
"url": None,
431+
"use_as_library_client": True,
432+
"api_key": "**********",
433+
"library_client_config_path": "tests/configuration/run.yaml",
434+
},
435+
"user_data_collection": {
436+
"feedback_enabled": False,
437+
"feedback_storage": None,
438+
"transcripts_enabled": False,
439+
"transcripts_storage": None,
440+
},
441+
"mcp_servers": [],
442+
"authentication": {
443+
"module": "noop",
444+
"skip_tls_verification": False,
445+
"k8s_ca_cert_path": None,
446+
"k8s_cluster_api": None,
447+
"jwk_config": None,
448+
},
449+
"customization": None,
450+
"inference": {
451+
"default_provider": "default_provider",
452+
"default_model": "default_model",
453+
},
454+
"database": {
455+
"sqlite": None,
456+
"postgres": {
457+
"host": "localhost",
458+
"port": 5432,
459+
"db": "lightspeed_stack",
460+
"user": "ls_user",
461+
"password": "**********",
462+
"ssl_mode": "require",
463+
"gss_encmode": "disable",
464+
"namespace": "lightspeed-stack",
465+
"ca_cert_path": None,
466+
},
467+
},
468+
"authorization": None,
469+
"conversation_cache": {
470+
"memory": None,
471+
"postgres": None,
472+
"sqlite": None,
473+
"type": None,
474+
},
475+
"byok_rag": [],
476+
"quota_handlers": {
477+
"sqlite": None,
478+
"postgres": None,
479+
"limiters": [
480+
{
481+
"initial_quota": 1,
482+
"name": "user_monthly_limits",
483+
"period": "2 seconds",
484+
"quota_increase": 10,
485+
"type": "user_limiter",
486+
},
487+
{
488+
"initial_quota": 2,
489+
"name": "cluster_monthly_limits",
490+
"period": "1 month",
491+
"quota_increase": 20,
492+
"type": "cluster_limiter",
493+
},
494+
],
495+
"scheduler": {"period": 10},
496+
"enable_token_history": True,
497+
},
498+
}
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
"""Unit tests for QuotaHandlersConfiguration model."""
2+
3+
import pytest
4+
5+
from models.config import QuotaHandlersConfiguration, QuotaSchedulerConfiguration
6+
7+
8+
def test_quota_handlers_configuration() -> None:
9+
"""Test the quota handlers configuration."""
10+
cfg = QuotaHandlersConfiguration(
11+
sqlite=None,
12+
postgres=None,
13+
limiters=[],
14+
scheduler=QuotaSchedulerConfiguration(period=10),
15+
enable_token_history=False,
16+
)
17+
assert cfg is not None
18+
assert cfg.sqlite is None
19+
assert cfg.postgres is None
20+
assert cfg.limiters == []
21+
assert cfg.scheduler is not None
22+
assert not cfg.enable_token_history
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
"""Unit tests for QuotaLimiterConfig model."""
2+
3+
import pytest
4+
5+
from models.config import QuotaLimiterConfiguration
6+
7+
8+
def test_quota_limiter_configuration() -> None:
9+
"""Test the default configuration."""
10+
cfg = QuotaLimiterConfiguration(
11+
type="cluster_limiter",
12+
name="cluster_monthly_limits",
13+
initial_quota=0,
14+
quota_increase=10,
15+
period="3 seconds",
16+
)
17+
assert cfg is not None
18+
assert cfg.type == "cluster_limiter"
19+
assert cfg.name == "cluster_monthly_limits"
20+
assert cfg.initial_quota == 0
21+
assert cfg.quota_increase == 10
22+
assert cfg.period == "3 seconds"
23+
24+
25+
def test_quota_limiter_configuration_improper_value_1() -> None:
26+
"""Test the default configuration."""
27+
with pytest.raises(ValueError, match="Input should be greater than or equal to 0"):
28+
_ = QuotaLimiterConfiguration(
29+
type="cluster_limiter",
30+
name="cluster_monthly_limits",
31+
initial_quota=-1,
32+
quota_increase=10,
33+
period="3 seconds",
34+
)
35+
36+
37+
def test_quota_limiter_configuration_improper_value_2() -> None:
38+
"""Test the default configuration."""
39+
with pytest.raises(ValueError, match="Input should be greater than or equal to 0"):
40+
_ = QuotaLimiterConfiguration(
41+
type="cluster_limiter",
42+
name="cluster_monthly_limits",
43+
initial_quota=1,
44+
quota_increase=-10,
45+
period="3 seconds",
46+
)
47+
48+
49+
def test_quota_limiter_configuration_improper_value_3() -> None:
50+
"""Test the default configuration."""
51+
with pytest.raises(
52+
ValueError, match="Input should be 'user_limiter' or 'cluster_limiter'"
53+
):
54+
_ = QuotaLimiterConfiguration(
55+
type="unknown_limiter",
56+
name="cluster_monthly_limits",
57+
initial_quota=1,
58+
quota_increase=10,
59+
period="3 seconds",
60+
)
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
"""Unit tests for QuotaSchedulerConfig model."""
2+
3+
import pytest
4+
5+
from models.config import QuotaSchedulerConfiguration
6+
7+
8+
def test_quota_scheduler_default_configuration() -> None:
9+
"""Test the default configuration."""
10+
cfg = QuotaSchedulerConfiguration()
11+
assert cfg is not None
12+
# default value
13+
assert cfg.period == 1
14+
15+
16+
def test_quota_scheduler_custom_configuration() -> None:
17+
"""Test the custom configuration."""
18+
cfg = QuotaSchedulerConfiguration(period=10)
19+
assert cfg is not None
20+
assert cfg.period == 10

0 commit comments

Comments
 (0)