Repository: qdrant/qdrant
File: tests/e2e_tests/test_snapshots_recovery.py (Python)

Code before:
        
        # Create collection and insert points
        collection_config = {"vectors": {"size": 4, "distance": "Dot"}}
        client.create_collection("test_collection", collection_config)
        for points_batch in client.generate_points(2):
            client.insert_points("test_collection", points_batch)

        # Create collection snapshot
        snapshot_name = client.create_snapshot("test_collection")
        assert snapshot_name, "Snapshot creation failed"
        
        # Download snapshot
        snapshot_content, snapshot_checksum = client.download_snapshot("test_collection", snapshot_name)
        assert len(snapshot_content) > 0, "Downloaded snapshot is empty"
        
        # Save snapshot to file for later use
        snapshot_file = tmp_path / "test_collection.snapshot"
        with open(snapshot_file, 'wb') as f:
            f.write(snapshot_content)
        
        # Test 1: Recover from URL
        snapshot_url = f"http://{client.host}:{client.port}/collections/test_collection/snapshots/{snapshot_name}"
        client.recover_snapshot_from_url("test_collection_recovered_1", snapshot_url, snapshot_checksum)
        
        # Verify recovered collection exists
        client.verify_collection_exists("test_collection_recovered_1")
        
        # Test 2: Upload snapshot as file
        client.upload_snapshot_file("test_collection_recovered_2", snapshot_content)
        
        # Verify second recovered collection exists
        client.verify_collection_exists("test_collection_recovered_2")
        
        # Test shard snapshots
        # Create shard snapshot
        shard_snapshot_name = client.create_shard_snapshot("test_collection", 0)
        assert shard_snapshot_name, "Shard snapshot creation failed"
        
        # Download shard snapshot
        shard_snapshot_content = client.download_shard_snapshot("test_collection", 0, shard_snapshot_name)
        assert len(shard_snapshot_content) > 0, "Downloaded shard snapshot is empty"
        
        # Save shard snapshot to file
        shard_snapshot_file = tmp_path / "test_collection_shard.snapshot"
        with open(shard_snapshot_file, 'wb') as f:
            f.write(shard_snapshot_content)
        
        # Test 3: Recover shard from URL
        shard_snapshot_url = f"http://{client.host}:{client.port}/collections/test_collection/shards/0/snapshots/{shard_snapshot_name}"
        client.recover_shard_snapshot_from_url("test_collection_recovered_1", 0, shard_snapshot_url)
        

Review comment (suggestion):
Maybe check `points_count`, `vectors_count`, `config`, etc returned by this? We can do the same for `recovered_2`

Code after:
        client.wait_for_server()
        
        # Create collection and insert points
        collection_config = {"vectors": {"size": 4, "distance": "Dot"}}
        client.create_collection("test_collection", collection_config)
        for points_batch in client.generate_points(2):
            client.insert_points("test_collection", points_batch, wait=True)
        initial_collection_config = client.get_collection_info_dict("test_collection")["result"]["config"]

        # Create collection snapshot
        snapshot_name = client.create_snapshot("test_collection")
        assert snapshot_name, "Snapshot creation failed"
        
        # Download snapshot
        snapshot_content, snapshot_checksum = client.download_snapshot("test_collection", snapshot_name)
        assert len(snapshot_content) > 0, "Downloaded snapshot is empty"
        
        # Save snapshot to file for later use
        snapshot_file = tmp_path / "test_collection.snapshot"
        with open(snapshot_file, 'wb') as f:
            f.write(snapshot_content)
        
        # Test 1: Recover from URL
        snapshot_url = f"http://{client.host}:{client.port}/collections/test_collection/snapshots/{snapshot_name}"
        client.recover_snapshot_from_url("test_collection_recovered_1", snapshot_url, snapshot_checksum)
        
        # Verify recovered collection exists
        resp = client.verify_collection_exists("test_collection_recovered_1")
        assert resp["result"].get("points_count") == 200, "Recovered collection does not have expected number of points"
        assert resp["result"].get("config") == initial_collection_config, "Recovered collection config does not match"

        # Test 2: Upload snapshot as file
        client.upload_snapshot_file("test_collection_recovered_2", snapshot_content)
        
        # Verify second recovered collection exists
        client.verify_collection_exists("test_collection_recovered_2")
        
        # Test shard snapshots
        # Create shard snapshot
        shard_snapshot_name = client.create_shard_snapshot("test_collection", 0)
        assert shard_snapshot_name, "Shard snapshot creation failed"
        
        # Download shard snapshot
        shard_snapshot_content = client.download_shard_snapshot("test_collection", 0, shard_snapshot_name)
        assert len(shard_snapshot_content) > 0, "Downloaded shard snapshot is empty"
        
        # Save shard snapshot to file
        shard_snapshot_file = tmp_path / "test_collection_shard.snapshot"
        with open(shard_snapshot_file, 'wb') as f:
            f.write(shard_snapshot_content)