do not skip late_bytes if none are posted

MegaAntiCheat · Jun 24, 2024 · a2bf12e · a2bf12e
1 parent e3d1332
commit a2bf12e
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/masterbase/lib.py b/masterbase/lib.py
@@ -428,17 +428,20 @@ def _close_session_with_demo(
                 "blob_name": demo_blob_name(session_id),
             },
         ).scalar_one()
-        with open(sink_path, "rb") as sink:
-            late = io.BytesIO(late_bytes)
-            head = io.BytesIO(sink.read(LATE_BYTES_START))
-            sink.seek(LATE_BYTES_END, os.SEEK_SET)
-            minio_client.put_object(
-                "demoblobs",
-                demo_blob_name(session_id),
-                data=cast(BinaryIO, ConcatStream(head, late, sink)),
-                length=size,
-                metadata={"has_late_bytes": str(bool(late_bytes))},
-            )
+        if late_bytes is not None:
+            with open(sink_path, "rb") as sink:
+                late = io.BytesIO(late_bytes)
+                head = io.BytesIO(sink.read(LATE_BYTES_START))
+                sink.seek(LATE_BYTES_END, os.SEEK_SET)
+                minio_client.put_object(
+                    "demoblobs",
+                    demo_blob_name(session_id),
+                    data=cast(BinaryIO, ConcatStream(head, late, sink)),
+                    length=size,
+                    metadata={"has_late_bytes": str(bool(late_bytes))},
+                )
+        else:
+            minio_client.fput_object("demoblobs", demo_blob_name(session_id), file_path=sink_path)
         conn.commit()
 
 

diff --git a/tests/test_api.py b/tests/test_api.py
@@ -104,7 +104,7 @@ def _send_demo_file(test_client: TestClient[Litestar], api_key: str, session_id:
 
 
 def test_demo_streaming(test_client: TestClient[Litestar], api_key: str) -> None:
-    """Test streaming a demo to the DB."""
+    """Test streaming a demo to the database with a header overwrite."""
     session_id = _open_mock_session(test_client, api_key).json()["session_id"]
     assert isinstance(session_id, int)
 
@@ -128,6 +128,24 @@ def test_demo_streaming(test_client: TestClient[Litestar], api_key: str) -> None
         assert demo_in == demo_out
 
 
+def test_demo_streaming_no_late(test_client: TestClient[Litestar], api_key: str) -> None:
+    """Test streaming a demo to the database without a header overwrite."""
+    session_id = _open_mock_session(test_client, api_key).json()["session_id"]
+    assert isinstance(session_id, int)
+
+    _send_demo_file(test_client, api_key, str(session_id))
+
+    close_session_response = test_client.get("/close_session", params={"api_key": api_key})
+    assert close_session_response.status_code == HTTP_200_OK
+
+    with test_client.stream("GET", "/demodata", params={"api_key": api_key, "session_id": session_id}) as demo_stream:
+        demo_out = demo_stream.read()
+
+    with open("tests/data/test_demo.dem", "rb") as f:
+        demo_in = f.read()
+        assert demo_in == demo_out
+
+
 def test_db_exports(test_client: TestClient[Litestar], api_key: str) -> None:
     """Test on-demand exports from the database."""
     session_id = str(_open_mock_session(test_client, api_key).json()["session_id"])