Test for RemoteArtifactSaver stage

test_prefetch_content.py

from django.test import TestCase
from pprint import pprint

from django.db import connection, reset_queries
from django.test.utils import override_settings
from pulpcore.plugin.models import Artifact, RemoteArtifact, ContentArtifact
from pulpcore.plugin.stages import (
    DeclarativeContent,
    DeclarativeArtifact,
    RemoteArtifactSaver,
)

from pulp_cookbook.app.models import CookbookPackageContent, CookbookRemote
from pulp_file.app.models import FileContent


class RemoteArtifactSaverTestCase(TestCase):
    def setUp(self):
        self.remote = CookbookRemote.objects.create(name="remote")
        self.remote2 = CookbookRemote.objects.create(name="remote2")
        self.remote3 = CookbookRemote.objects.create(name="remote3")

        self.cs = []
        self.d_cs = []
        self.ras = []
        for n in range(100):
            c_name = f"c{n:04d}"
            a_url = f"http://{c_name}_remote"
            # Test different content types: "even content" is FileContent, "odd
            # content" is CookbookPackageContent
            if n % 2:
                c = CookbookPackageContent.objects.create(
                    name=c_name, version="1.0.0", dependencies={}
                )
                rel_path = c.relative_path()
            else:
                rel_path = c_name
                c = FileContent.objects.create(relative_path=rel_path)
            ca = ContentArtifact.objects.create(
                artifact=None, content=c, relative_path=rel_path
            )

            # Every third Content.artifact has a matching RemoteArtifact
            if n % 3 == 0:
                ra = RemoteArtifact.objects.create(
                    url=a_url, remote=self.remote, content_artifact=ca
                )
                self.ras.append(ra)
            # Every fifth Content.artifact has an unsuitable RemoteArtifact (different remote)
            if n % 5 == 0:
                ra = RemoteArtifact.objects.create(
                    url=a_url + "2", remote=self.remote2, content_artifact=ca
                )
                self.ras.append(ra)

            d_a = DeclarativeArtifact(
                artifact=Artifact(),
                url=a_url,
                relative_path=rel_path,
                remote=self.remote,
            )

            d_as = [d_a]

            # Every seventh Content has an additional artifact. This Artifact
            # actually exists in the DB (should make no difference at all). For
            # every second of these, there is a matching RemoteArtifact using
            # remote3.
            if n % 7 == 0:
                Artifact
                a = Artifact.objects.create(
                    size=n, sha256=str(n), sha384=str(n), sha512=str(n)
                )
                ca = ContentArtifact.objects.create(
                    artifact=a, content=c, relative_path="3/" + rel_path
                )
                d_a2 = DeclarativeArtifact(
                    artifact=a,
                    url=a_url + "3",
                    relative_path="3/" + rel_path,
                    remote=self.remote3,
                )
                d_as.append(d_a2)
                if n % 2 == 0:
                    ra = RemoteArtifact.objects.create(
                        url=a_url + "3", remote=self.remote3, content_artifact=ca
                    )
                    self.ras.append(ra)

            d_c = DeclarativeContent(c, d_artifacts=d_as)

            self.cs.append(c)
            self.d_cs.append(d_c)

    @override_settings(DEBUG=True)
    def test_needed_remote_artifacts(self):
        stage = RemoteArtifactSaver()
        reset_queries()
        needed_urls = sorted(
            [ra.url for ra in stage._needed_remote_artifacts(self.d_cs)]
        )
        print(f"Number of DB queries: {len(connection.queries)}")
        pprint(connection.queries)
        self.assertEqual(2, len(connection.queries))
        expected_needed_urls = []
        for n in range(100):
            if n % 3:
                expected_needed_urls.append(f"http://c{n:04d}_remote")
            if n % 7 == 0 and n % 2:
                expected_needed_urls.append(f"http://c{n:04d}_remote3")
        self.assertListEqual(needed_urls, expected_needed_urls)